From 1624385cee716f8eb2aa752cba4fb7f95dbfff86 Mon Sep 17 00:00:00 2001 From: suguanYang Date: Fri, 11 Sep 2026 19:49:06 +0800 Subject: [PATCH] fix(demo): deduplicate materialization and parallelize result uploads --- ...5a_add_demo_materialization_claim_state.py | 52 ++++ .../app/services/demo/source_materializer.py | 222 ++++++++++++------ .../contract/test_demo_documents_contract.py | 67 ++++-- .../test_page_memory_retrieval_contract.py | 174 +++++++++++++- .../shared/core/config/storage.py | 16 ++ .../models/database/demo_materialization.py | 10 +- .../services/retrieval/publication_content.py | 13 +- .../services/retrieval/publication_models.py | 2 + .../services/retrieval/publication_service.py | 81 ++++--- .../shared/services/storage/result_storage.py | 81 ++++++- .../tests/test_storage_config_contract.py | 7 + 11 files changed, 576 insertions(+), 149 deletions(-) create mode 100644 apps/api/alembic/versions/0b1c2d3e4f5a_add_demo_materialization_claim_state.py diff --git a/apps/api/alembic/versions/0b1c2d3e4f5a_add_demo_materialization_claim_state.py b/apps/api/alembic/versions/0b1c2d3e4f5a_add_demo_materialization_claim_state.py new file mode 100644 index 000000000..81f676d02 --- /dev/null +++ b/apps/api/alembic/versions/0b1c2d3e4f5a_add_demo_materialization_claim_state.py @@ -0,0 +1,52 @@ +"""Add demo materialization claim state.""" + +from __future__ import annotations + +from collections.abc import Sequence + +from alembic import op +import sqlalchemy as sa + + +revision: str = "0b1c2d3e4f5a" +down_revision: str | Sequence[str] | None = "e4f5a6b7c8d9" +branch_labels: Sequence[str] | None = None +depends_on: Sequence[str] | None = None + + +def upgrade() -> None: + op.add_column( + "demo_materializations", + sa.Column("status", sa.String(length=32), nullable=True), + ) + op.add_column( + "demo_materializations", + sa.Column("claimed_at", sa.DateTime(), nullable=True), + ) + op.execute( + "UPDATE demo_materializations SET status = 'ready' WHERE status IS NULL" + ) + op.alter_column( + "demo_materializations", + "status", + existing_type=sa.String(length=32), + nullable=False, + server_default="ready", + ) + op.alter_column( + "demo_materializations", + "document_id", + existing_type=sa.String(length=36), + nullable=True, + ) + + +def downgrade() -> None: + op.alter_column( + "demo_materializations", + "document_id", + existing_type=sa.String(length=36), + nullable=False, + ) + op.drop_column("demo_materializations", "claimed_at") + op.drop_column("demo_materializations", "status") diff --git a/apps/api/app/services/demo/source_materializer.py b/apps/api/app/services/demo/source_materializer.py index df9e2f0c8..0f0bef48f 100644 --- a/apps/api/app/services/demo/source_materializer.py +++ b/apps/api/app/services/demo/source_materializer.py @@ -4,23 +4,25 @@ import shutil import tempfile +from collections.abc import Iterable from dataclasses import dataclass -from datetime import datetime, timezone +from datetime import datetime, timedelta, timezone from hashlib import blake2b from pathlib import Path from uuid import uuid4 from app.services.demo.source_catalog import DemoSourceCatalog, DemoSourceDefinition -from sqlalchemy import func, select +from sqlalchemy.exc import IntegrityError +from sqlalchemy import delete, func, select from sqlalchemy.ext.asyncio import AsyncSession -from shared.core.exceptions.domain_exceptions import ValidationException +from shared.core.exceptions.domain_exceptions import ConflictException, ValidationException from shared.models.database.demo_materialization import DemoMaterialization -from shared.models.database.document import Document from shared.models.database.job import Job from shared.models.database.job_result import JobResult from shared.services.retrieval.cache_service import invalidate_retrieval_cache_namespaces from shared.services.retrieval.publication_service import RetrievalPublicationService +from shared.services.retrieval.publication_models import DocumentPublicationScope from shared.services.storage.result_storage import get_result_storage @@ -73,17 +75,39 @@ async def materialize_sources( self._catalog.require_source(demo_source_id) for demo_source_id in selected_demo_source_ids ] - results: list[MaterializedDemoSource] = [] - for source in selected_sources: - result = await self._materialize_source( + try: + claims = await self._claim_sources( db, user_id=user_id, namespace=namespace, - source=source, + sources=selected_sources, ) - results.append(result) - + except IntegrityError as error: + await db.rollback() + raise ConflictException( + user_message="This demo source is currently being materialized.", + reason="ABORTED", + resource="Demo materialization", + ) from error await db.commit() + results: list[MaterializedDemoSource] = [] + for source in selected_sources: + try: + result = await self._materialize_source( + db, + user_id=user_id, + namespace=namespace, + source=source, + claim=claims[source.demo_source_id], + ) + results.append(result) + except Exception: + await db.rollback() + await self._release_claims( + db, + claims=claims.values(), + ) + raise await invalidate_retrieval_cache_namespaces( user_id=user_id, namespaces=[namespace], @@ -97,29 +121,8 @@ async def _materialize_source( user_id: str, namespace: str, source: DemoSourceDefinition, + claim: DemoMaterialization, ) -> MaterializedDemoSource: - await _lock_materialization_scope( - db, - user_id=user_id, - namespace=namespace, - demo_source_id=source.demo_source_id, - ) - existing = await self._get_existing_materialization( - db, - user_id=user_id, - namespace=namespace, - demo_source_id=source.demo_source_id, - ) - if existing is not None and await self._is_active_document( - db, - document_id=existing.document_id, - ): - return _materialized_source_payload( - source=source, - document_id=existing.document_id, - status="existing", - ) - document_id = f"doc_{uuid4().hex[:12]}" job_id = f"job_demo_{uuid4().hex[:12]}" job_result_id = str(uuid4()) @@ -169,12 +172,13 @@ async def _materialize_source( ) await db.flush() chunks = self._catalog.publication_chunks(source) - await db.run_sync( + published_state = await db.run_sync( lambda sync_db: self._publication_service.publish_document_state( sync_db, job_id=job_id, job_result_id=job_result_id, chunks=[dict(chunk) for chunk in chunks], + update_namespace_snapshot=False, ) ) await db.run_sync( @@ -186,58 +190,114 @@ async def _materialize_source( ) await db.flush() - if existing is None: - db.add( - DemoMaterialization( + if published_state is None or published_state.document_id != document_id: + raise RuntimeError("Demo publication did not create its requested document") + if published_state.manifest_payload is None: + raise RuntimeError("Demo publication did not create a serving manifest") + manifest_payload = published_state.manifest_payload + await db.run_sync( + lambda sync_db: self._publication_service.update_namespace_snapshot( + sync_db, + scope=DocumentPublicationScope( user_id=user_id, namespace=namespace, - demo_source_id=source.demo_source_id, document_id=document_id, - created_at=timestamp, - updated_at=timestamp, - ) + job_result_id=job_result_id, + source_file_name=source.title, + ), + manifest_payload=manifest_payload, ) - else: - existing.document_id = document_id - existing.updated_at = timestamp + ) + await db.commit() + + claim.document_id = document_id + claim.status = "ready" + claim.claimed_at = None + claim.updated_at = timestamp await db.flush() + await db.commit() return _materialized_source_payload( source=source, document_id=document_id, status="created", ) - async def _get_existing_materialization( + async def _claim_sources( self, db: AsyncSession, *, user_id: str, namespace: str, - demo_source_id: str, - ) -> DemoMaterialization | None: - result = await db.execute( - select(DemoMaterialization) - .where(DemoMaterialization.user_id == user_id) - .where(DemoMaterialization.namespace == namespace) - .where(DemoMaterialization.demo_source_id == demo_source_id) - .with_for_update() - .limit(1) - ) - return result.scalar_one_or_none() + sources: list[DemoSourceDefinition], + ) -> dict[str, DemoMaterialization]: + now = _utc_now() + claims: dict[str, DemoMaterialization] = {} + for source in sorted(sources, key=lambda item: item.demo_source_id): + lock_acquired = await db.scalar( + select( + func.pg_try_advisory_xact_lock( + _materialization_lock_id( + user_id=user_id, + namespace=namespace, + demo_source_id=source.demo_source_id, + ) + ) + ) + ) + if not lock_acquired: + raise ConflictException( + user_message="This demo source is currently being materialized.", + reason="ABORTED", + resource="Demo materialization", + resource_id=source.demo_source_id, + ) + result = await db.execute( + select(DemoMaterialization) + .where(DemoMaterialization.user_id == user_id) + .where(DemoMaterialization.namespace == namespace) + .where(DemoMaterialization.demo_source_id == source.demo_source_id) + ) + existing = result.scalar_one_or_none() + if existing is not None: + if not _is_stale_claim(existing, now=now): + raise _materialization_conflict(existing, source) + claim = existing + claim.status = "materializing" + claim.document_id = None + claim.claimed_at = now + claim.updated_at = now + else: + claim = DemoMaterialization( + user_id=user_id, + namespace=namespace, + demo_source_id=source.demo_source_id, + status="materializing", + document_id=None, + claimed_at=now, + created_at=now, + updated_at=now, + ) + db.add(claim) + claims[source.demo_source_id] = claim + await db.flush() + return claims - async def _is_active_document( + async def _release_claims( self, db: AsyncSession, *, - document_id: str, - ) -> bool: - result = await db.execute( - select(Document.document_id) - .where(Document.document_id == document_id) - .where(Document.status == "active") - .limit(1) + claims: Iterable[DemoMaterialization], + ) -> None: + claim_ids = [claim.id for claim in claims] + if not claim_ids: + return + await db.execute( + delete(DemoMaterialization).where( + DemoMaterialization.status == "materializing", + DemoMaterialization.id.in_(claim_ids), + ) ) - return result.scalar_one_or_none() is not None + await db.commit() def _deduplicate_source_ids(demo_source_ids: list[str]) -> list[str]: @@ -252,19 +312,12 @@ def _deduplicate_source_ids(demo_source_ids: list[str]) -> list[str]: return selected -async def _lock_materialization_scope( - db: AsyncSession, - *, - user_id: str, - namespace: str, - demo_source_id: str, -) -> None: - lock_id = _materialization_lock_id( - user_id=user_id, - namespace=namespace, - demo_source_id=demo_source_id, +def _is_stale_claim(materialization: DemoMaterialization, *, now: datetime) -> bool: + return ( + materialization.status == "materializing" + and materialization.claimed_at is not None + and materialization.claimed_at < now - timedelta(minutes=10) ) - await db.execute(select(func.pg_advisory_xact_lock(lock_id))) def _materialization_lock_id( @@ -278,6 +331,23 @@ def _materialization_lock_id( return int.from_bytes(digest, byteorder="big", signed=True) +def _materialization_conflict( + materialization: DemoMaterialization, + source: DemoSourceDefinition, +) -> ConflictException: + reason = "ALREADY_EXISTS" if materialization.status == "ready" else "ABORTED" + return ConflictException( + user_message=( + "This demo source has already been materialized." + if reason == "ALREADY_EXISTS" + else "This demo source is currently being materialized." + ), + reason=reason, + resource="Demo materialization", + resource_id=source.demo_source_id, + ) + + def _materialized_source_payload( *, source: DemoSourceDefinition, diff --git a/apps/api/tests/contract/test_demo_documents_contract.py b/apps/api/tests/contract/test_demo_documents_contract.py index bfe13e105..f2eab17e2 100644 --- a/apps/api/tests/contract/test_demo_documents_contract.py +++ b/apps/api/tests/contract/test_demo_documents_contract.py @@ -365,7 +365,7 @@ async def test_should_materialize_demo_source_without_parse_or_credit_charge( assert empty_cached_response.status_code == 200 assert first_response.status_code == 200 - assert retry_response.status_code == 200 + assert retry_response.status_code == 409 assert retrieval_response.status_code == 200 assert document_chunks_response.status_code == 200 @@ -373,12 +373,10 @@ async def test_should_materialize_demo_source_without_parse_or_credit_charge( assert cast(list[dict[str, Any]], empty_cached_body["results"]) == [] first_source = cast(dict[str, Any], first_response.json()["sources"][0]) - retry_source = cast(dict[str, Any], retry_response.json()["sources"][0]) document_id = str(first_source["document_id"]) assert first_source["status"] == "created" - assert retry_source["status"] == "existing" - assert retry_source["document_id"] == document_id + assert retry_response.json()["error"]["details"]["reason"] == "ALREADY_EXISTS" materialization_rows = await ContractDatabase.fetch_all( """ @@ -517,7 +515,39 @@ async def test_should_materialize_each_normalized_demo_source_once_per_request( @pytest.mark.asyncio -async def test_should_serialize_concurrent_first_demo_materialization( +async def test_should_reject_sequential_duplicate_demo_materialization( + developer_api_client_factory: Callable[ + [], AbstractAsyncContextManager[AsyncClient] + ], + monkeypatch: MonkeyPatch, +) -> None: + fake_result_storage = FakeResultStorage() + + async with developer_api_client_factory() as api_client: + import app.services.demo.source_materializer as source_materializer_module + + monkeypatch.setattr( + source_materializer_module, + "get_result_storage", + lambda: fake_result_storage, + ) + first_response = await api_client.post( + "/api/v1/demo/materializations", + json={"namespace": "contract-demo-duplicate", "demo_source_ids": [DEMO_SOURCE_ID]}, + ) + second_response = await api_client.post( + "/api/v1/demo/materializations", + json={"namespace": "contract-demo-duplicate", "demo_source_ids": [DEMO_SOURCE_ID]}, + ) + + assert first_response.status_code == 200 + assert second_response.status_code == 409 + assert second_response.json()["error"]["details"]["reason"] == "ALREADY_EXISTS" + assert len(fake_result_storage.raw_files_by_job_id) == 1 + + +@pytest.mark.asyncio +async def test_should_reject_concurrent_duplicate_demo_materialization( developer_api_client_factory: Callable[ [], AbstractAsyncContextManager[AsyncClient] ], @@ -550,19 +580,12 @@ async def test_should_serialize_concurrent_first_demo_materialization( ), ) - assert first_response.status_code == 200 - assert second_response.status_code == 200 - - first_source = cast(dict[str, Any], first_response.json()["sources"][0]) - second_source = cast(dict[str, Any], second_response.json()["sources"][0]) - document_ids = { - str(first_source["document_id"]), - str(second_source["document_id"]), - } - statuses = {str(first_source["status"]), str(second_source["status"])} - - assert len(document_ids) == 1 - assert statuses == {"created", "existing"} + statuses = {first_response.status_code, second_response.status_code} + assert statuses == {200, 409} + conflict_response = ( + first_response if first_response.status_code == 409 else second_response + ) + assert conflict_response.json()["error"]["details"]["reason"] == "ABORTED" materialization_rows = await ContractDatabase.fetch_all( """ @@ -585,12 +608,8 @@ async def test_should_serialize_concurrent_first_demo_materialization( {"demo_source_id": DEMO_SOURCE_ID}, ) - assert materialization_rows == [ - { - "demo_source_id": DEMO_SOURCE_ID, - "document_id": next(iter(document_ids)), - } - ] + assert len(materialization_rows) == 1 + assert materialization_rows[0]["document_id"] is not None assert len(job_rows) == 1 diff --git a/apps/worker/tests/contract/test_page_memory_retrieval_contract.py b/apps/worker/tests/contract/test_page_memory_retrieval_contract.py index 1319db3fb..e9346c7ed 100644 --- a/apps/worker/tests/contract/test_page_memory_retrieval_contract.py +++ b/apps/worker/tests/contract/test_page_memory_retrieval_contract.py @@ -2,6 +2,8 @@ import os import shutil +import threading +import time from pathlib import Path import pytest @@ -29,6 +31,7 @@ from shared.services.retrieval.execution.reference_resolver import ( # noqa: E402 resolve_workflow_references, ) +from shared.core.exceptions.domain_exceptions import StorageServiceException # noqa: E402 from shared.services.storage.result_storage import JobResultStorage # noqa: E402 from shared.services.storage.page_pdf_crop import crop_source_pdf_pages # noqa: E402 @@ -267,7 +270,9 @@ class FakeStorageAdapter: def __init__(self) -> None: self.uploaded_keys: list[str] = [] - def upload_file(self, local_path: str, key: str, bucket: str | None = None): + def upload_file( + self, local_path: str, key: str, bucket: str | None = None + ) -> dict[str, str]: del local_path, bucket self.uploaded_keys.append(key) return {"key": key} @@ -322,6 +327,173 @@ def generate_presigned_url(self, *args, **kwargs) -> str: assert "results/job-1/debug.csv" not in adapter.uploaded_keys +def test_result_storage_uploads_raw_files_with_bounded_concurrency(tmp_path) -> None: + class ConcurrentStorageAdapter: + def __init__(self) -> None: + self.uploaded_keys: list[str] = [] + self.active_uploads = 0 + self.maximum_active_uploads = 0 + self.lock = threading.Lock() + + def upload_file( + self, local_path: str, key: str, bucket: str | None = None + ) -> dict[str, str]: + del local_path, bucket + with self.lock: + self.active_uploads += 1 + self.maximum_active_uploads = max( + self.maximum_active_uploads, + self.active_uploads, + ) + time.sleep(0.01) + with self.lock: + self.uploaded_keys.append(key) + self.active_uploads -= 1 + return {"key": key} + + def generate_presigned_url(self, *args, **kwargs) -> str: + del args, kwargs + return "https://assets.example.test/file" + + result_dir = tmp_path / "result" + result_dir.mkdir() + for index in range(8): + (result_dir / f"asset-{index}.bin").write_bytes(b"asset") + zip_path = tmp_path / "result.zip" + zip_path.write_bytes(b"zip") + adapter = ConcurrentStorageAdapter() + storage = JobResultStorage( + results_bucket="test-results", + storage_adapter=adapter, # type: ignore[arg-type] + upload_concurrency=3, + ) + + bundle = storage.upload( + job_id="job-concurrent", + result_dir=str(result_dir), + zip_file_path=str(zip_path), + ) + + assert len(bundle.raw_files) == 8 + assert adapter.maximum_active_uploads <= 3 + assert adapter.maximum_active_uploads > 1 + + +def test_result_storage_bounds_uploads_across_concurrent_tasks(tmp_path) -> None: + class ProcessStorageAdapter: + def __init__(self) -> None: + self.active_uploads = 0 + self.maximum_active_uploads = 0 + self.lock = threading.Lock() + + def upload_file( + self, local_path: str, key: str, bucket: str | None = None + ) -> dict[str, str]: + del local_path, bucket + if key.endswith(".zip"): + return {"key": key} + with self.lock: + self.active_uploads += 1 + self.maximum_active_uploads = max( + self.maximum_active_uploads, + self.active_uploads, + ) + time.sleep(0.02) + with self.lock: + self.active_uploads -= 1 + return {"key": key} + + def generate_presigned_url(self, *args, **kwargs) -> str: + del args, kwargs + return "https://assets.example.test/file" + + def create_result_directory(name: str) -> tuple[Path, Path]: + result_dir = tmp_path / name + result_dir.mkdir() + for index in range(8): + (result_dir / f"asset-{index}.bin").write_bytes(b"asset") + zip_path = tmp_path / f"{name}.zip" + zip_path.write_bytes(b"zip") + return result_dir, zip_path + + adapter = ProcessStorageAdapter() + first_result_dir, first_zip_path = create_result_directory("first-result") + second_result_dir, second_zip_path = create_result_directory("second-result") + first_storage = JobResultStorage( + storage_adapter=adapter, # type: ignore[arg-type] + upload_concurrency=3, + ) + second_storage = JobResultStorage( + storage_adapter=adapter, # type: ignore[arg-type] + upload_concurrency=3, + ) + failures: list[Exception] = [] + + def upload_result( + storage: JobResultStorage, + job_id: str, + result_dir: Path, + zip_path: Path, + ) -> None: + try: + storage.upload( + job_id=job_id, + result_dir=str(result_dir), + zip_file_path=str(zip_path), + ) + except Exception as exc: + failures.append(exc) + + first_thread = threading.Thread( + target=upload_result, + args=(first_storage, "first-job", first_result_dir, first_zip_path), + ) + second_thread = threading.Thread( + target=upload_result, + args=(second_storage, "second-job", second_result_dir, second_zip_path), + ) + + first_thread.start() + second_thread.start() + first_thread.join() + second_thread.join() + + assert failures == [] + assert adapter.maximum_active_uploads == 3 + + +def test_result_storage_upload_propagates_raw_file_failure(tmp_path) -> None: + class FailingStorageAdapter: + def upload_file(self, local_path: str, key: str, bucket: str | None = None): + del local_path, bucket + if key.endswith("asset-2.bin"): + raise RuntimeError("upload failed") + return {"key": key} + + def generate_presigned_url(self, *args, **kwargs) -> str: + del args, kwargs + return "https://assets.example.test/file" + + result_dir = tmp_path / "result" + result_dir.mkdir() + for index in range(4): + (result_dir / f"asset-{index}.bin").write_bytes(b"asset") + zip_path = tmp_path / "result.zip" + zip_path.write_bytes(b"zip") + storage = JobResultStorage( + results_bucket="test-results", + storage_adapter=FailingStorageAdapter(), # type: ignore[arg-type] + upload_concurrency=2, + ) + + with pytest.raises(StorageServiceException, match="Storage upload failed|upload failed"): + storage.upload( + job_id="job-failure", + result_dir=str(result_dir), + zip_file_path=str(zip_path), + ) + + def test_crop_source_pdf_pages_uploads_and_reuses_page_pdf_cache(tmp_path) -> None: from pypdf import PdfReader, PdfWriter diff --git a/packages/shared-python/shared/core/config/storage.py b/packages/shared-python/shared/core/config/storage.py index e04d69d64..930ae06cd 100644 --- a/packages/shared-python/shared/core/config/storage.py +++ b/packages/shared-python/shared/core/config/storage.py @@ -52,6 +52,12 @@ class StorageConfig(BaseModel): S3_ADDRESSING_STYLE: str = Field( default="auto", description="S3 addressing style: auto, path, or virtual" ) + S3_MAX_POOL_CONNECTIONS: int = Field( + default=20, + ge=1, + le=256, + description="Maximum connections retained by the shared S3 client pool.", + ) # OSS-only configuration. OSS_ENDPOINT: str = Field( @@ -96,6 +102,15 @@ class StorageConfig(BaseModel): le=10, description="Maximum concurrent MinerU API calls for shard parsing.", ) + RESULT_UPLOAD_CONCURRENCY: int = Field( + default=20, + ge=1, + le=64, + description=( + "Maximum concurrent raw result-file uploads per materialization. " + "ZIP bundles remain single-object uploads." + ), + ) SUPPORTED_EXTENSIONS: str = Field( default=".doc,.docx,.pdf,.txt,.xls,.xlsx,.pptx,.jpg,.jpeg,.png,.md,.html,.htm", description="Supported file extensions", @@ -146,6 +161,7 @@ def get_s3_client(self) -> BaseClient: # Configure retries. config_kwargs["retries"] = {"max_attempts": 5, "mode": "standard"} + config_kwargs["max_pool_connections"] = self.S3_MAX_POOL_CONNECTIONS config = Config(**config_kwargs) if config_kwargs else None diff --git a/packages/shared-python/shared/models/database/demo_materialization.py b/packages/shared-python/shared/models/database/demo_materialization.py index d4cc1a7a1..32bbc950b 100644 --- a/packages/shared-python/shared/models/database/demo_materialization.py +++ b/packages/shared-python/shared/models/database/demo_materialization.py @@ -29,10 +29,16 @@ class DemoMaterialization(Base): String(255), nullable=False, default="default" ) demo_source_id: Mapped[str] = mapped_column(String(128), nullable=False) - document_id: Mapped[str] = mapped_column( + status: Mapped[str] = mapped_column( + String(32), nullable=False, default="ready" + ) + document_id: Mapped[str | None] = mapped_column( String(36), ForeignKey("documents.document_id", ondelete="CASCADE"), - nullable=False, + nullable=True, + ) + claimed_at: Mapped[datetime | None] = mapped_column( + DateTime, nullable=True ) created_at: Mapped[datetime] = mapped_column( DateTime, default=utc_now_naive, nullable=False diff --git a/packages/shared-python/shared/services/retrieval/publication_content.py b/packages/shared-python/shared/services/retrieval/publication_content.py index 7ce75b808..90ab74e79 100644 --- a/packages/shared-python/shared/services/retrieval/publication_content.py +++ b/packages/shared-python/shared/services/retrieval/publication_content.py @@ -13,9 +13,6 @@ DocumentSection, ) from shared.services.retrieval.map_unit_index import replace_document_map_units -from shared.services.retrieval.namespace_map_snapshot import ( - patch_namespace_map_snapshot, -) from shared.services.retrieval.publication_models import DocumentPublicationScope from shared.services.retrieval.serving_manifest import persist_revision_serving_state from shared.services.retrieval.search.lexical_text import ( @@ -63,7 +60,7 @@ def replace_document_revision_content( scope: DocumentPublicationScope, chunks: list[dict[str, Any]], section_summaries: dict[str, str] | None = None, -) -> None: +) -> dict[str, Any]: """Replace retrieval sections and chunks for one published document revision.""" _delete_existing_revision_content(db, scope=scope) section_publisher = DocumentSectionPublisher( @@ -71,6 +68,7 @@ def replace_document_revision_content( scope=scope, section_summaries=section_summaries, ) + prepared_chunks: list[tuple[int, dict[str, Any], dict[str, Any], str | None, DocumentSection]] = [] for index, chunk in enumerate(chunks): safe_chunk = cast(dict[str, Any], remove_nul_characters(chunk)) chunk_metadata = _get_chunk_metadata(safe_chunk) @@ -83,6 +81,9 @@ def replace_document_revision_content( source_file_name=scope.source_file_name, ) section = section_publisher.ensure_section(section_path) + prepared_chunks.append((index, safe_chunk, chunk_metadata, source_path, section)) + db.flush() + for index, safe_chunk, chunk_metadata, source_path, section in prepared_chunks: db.add( _build_document_chunk( chunk=safe_chunk, @@ -97,7 +98,7 @@ def replace_document_revision_content( replace_document_map_units(db, scope=scope) db.flush() manifest_payload = persist_revision_serving_state(db, scope=scope) - patch_namespace_map_snapshot(db, scope=scope, manifest_payload=manifest_payload) + return manifest_payload class DocumentSectionPublisher: @@ -128,6 +129,7 @@ def ensure_section(self, section_path: str) -> DocumentSection: continue ancestor_section = DocumentSection( + section_id=f"sec_{uuid4().hex[:12]}", user_id=self._scope.user_id, namespace=self._scope.namespace, document_id=self._scope.document_id, @@ -141,7 +143,6 @@ def ensure_section(self, section_path: str) -> DocumentSection: summary=self._section_summaries.get(ancestor_path) or None, ) self._db.add(ancestor_section) - self._db.flush() self._sections_by_path[ancestor_path] = ancestor_section return self._sections_by_path[section_path] diff --git a/packages/shared-python/shared/services/retrieval/publication_models.py b/packages/shared-python/shared/services/retrieval/publication_models.py index 8c653212c..0fdea6b11 100644 --- a/packages/shared-python/shared/services/retrieval/publication_models.py +++ b/packages/shared-python/shared/services/retrieval/publication_models.py @@ -1,6 +1,7 @@ from __future__ import annotations from dataclasses import dataclass +from typing import Any @dataclass(frozen=True) @@ -15,6 +16,7 @@ class PublishedDocumentState: namespace: str document_id: str | None skipped_all_duplicate: bool = False + manifest_payload: dict[str, Any] | None = None @dataclass(frozen=True) diff --git a/packages/shared-python/shared/services/retrieval/publication_service.py b/packages/shared-python/shared/services/retrieval/publication_service.py index 7ee5dec2b..3cdd8ddc0 100644 --- a/packages/shared-python/shared/services/retrieval/publication_service.py +++ b/packages/shared-python/shared/services/retrieval/publication_service.py @@ -23,6 +23,7 @@ from shared.models.schemas.retrieval_namespace import normalize_retrieval_namespace from shared.services.retrieval.graph.service import DocumentGraphService, GraphScope from shared.services.retrieval.namespace_map_snapshot import ( + patch_namespace_map_snapshot, remove_document_from_namespace_map_snapshot, ) from shared.services.retrieval.publication_content import ( @@ -81,6 +82,7 @@ def publish_document_state( job_result_id: str, chunks: list[dict[str, Any]], section_summaries: dict[str, str] | None = None, + update_namespace_snapshot: bool = True, ) -> PublishedDocumentState | None: job = db.execute(select(Job).where(Job.job_id == job_id)).scalar_one_or_none() if not job: @@ -93,6 +95,7 @@ def publish_document_state( job_result_id=job_result_id, chunks=chunks, section_summaries=section_summaries, + update_namespace_snapshot=update_namespace_snapshot, ) def _publish_document_state_for_job( @@ -103,6 +106,7 @@ def _publish_document_state_for_job( job_result_id: str, chunks: list[dict[str, Any]], section_summaries: dict[str, str] | None = None, + update_namespace_snapshot: bool = True, ) -> PublishedDocumentState | None: job_metadata = job.job_metadata or {} @@ -136,16 +140,6 @@ def _publish_document_state_for_job( Document.document_id == str(document_id) ) ).scalar_one_or_none() - namespaces_to_lock = {namespace} - if existing_namespace: - namespaces_to_lock.add(str(existing_namespace)) - for namespace_to_lock in sorted(namespaces_to_lock): - lock_namespace_generation( - db, - user_id=str(job.user_id), - namespace=namespace_to_lock, - ) - document = self._upsert_document_revision( db, job=job, @@ -172,7 +166,7 @@ def _publish_document_state_for_job( job_result_id=job_result_id, source_file_name=str(source_file_name) if source_file_name else None, ) - replace_document_revision_content( + manifest_payload = replace_document_revision_content( db, scope=scope, chunks=deduped_chunks, @@ -180,32 +174,61 @@ def _publish_document_state_for_job( ) db.flush() - if existing_namespace and str(existing_namespace) != scope.namespace: - remove_document_from_namespace_map_snapshot( + if update_namespace_snapshot: + self.update_namespace_snapshot( db, - user_id=scope.user_id, - namespace=str(existing_namespace), - document_id=document.document_id, + scope=scope, + manifest_payload=manifest_payload, + previous_namespace=( + str(existing_namespace) + if existing_namespace and str(existing_namespace) != scope.namespace + else None + ), ) - # A namespace move mutates both namespace snapshots. Advance the - # old namespace generation as well so request-scoped/process-local - # snapshot caches cannot reuse the pre-move generation. - advance_namespace_generation( - db, - user_id=scope.user_id, - namespace=str(existing_namespace), - ) - advance_namespace_generation( - db, - user_id=scope.user_id, - namespace=scope.namespace, - ) return PublishedDocumentState( user_id=str(job.user_id), namespace=namespace, document_id=document.document_id, + manifest_payload=manifest_payload, ) + def update_namespace_snapshot( + self, + db: Session, + *, + scope: DocumentPublicationScope, + manifest_payload: dict[str, Any], + previous_namespace: str | None = None, + ) -> None: + """Patch one namespace snapshot while holding only its short lock.""" + namespaces = [scope.namespace] + if previous_namespace: + namespaces.insert(0, previous_namespace) + for namespace in namespaces: + lock_namespace_generation( + db, + user_id=scope.user_id, + namespace=namespace, + ) + if namespace == scope.namespace: + patch_namespace_map_snapshot( + db, + scope=scope, + manifest_payload=manifest_payload, + ) + else: + remove_document_from_namespace_map_snapshot( + db, + user_id=scope.user_id, + namespace=namespace, + document_id=scope.document_id, + ) + advance_namespace_generation( + db, + user_id=scope.user_id, + namespace=namespace, + ) + def _upsert_document_revision( self, db: Session, diff --git a/packages/shared-python/shared/services/storage/result_storage.py b/packages/shared-python/shared/services/storage/result_storage.py index 51d3ab3ce..9a5643fb1 100644 --- a/packages/shared-python/shared/services/storage/result_storage.py +++ b/packages/shared-python/shared/services/storage/result_storage.py @@ -1,13 +1,16 @@ from __future__ import annotations import os +import threading from collections.abc import Iterator +from concurrent.futures import ThreadPoolExecutor from dataclasses import dataclass from pathlib import Path from typing import Protocol from loguru import logger +from shared.core.config import settings from shared.services.storage.job_file_storage import JobFileStorage from shared.services.storage.storage_adapter import StorageAdapter @@ -15,6 +18,9 @@ _EXCLUDED_DIR_NAMES = {"tmp", "temp", "__pycache__"} _CLIENT_ARTIFACT_DIRS = {"images", "tables", "page_pdfs", "page_citation_assets"} _INTERNAL_RAW_FILES = {"source.pdf"} +_MAX_RESULT_UPLOAD_CONCURRENCY = 64 +_upload_executor_lock = threading.Lock() +_upload_executors: dict[int, ThreadPoolExecutor] = {} @dataclass(frozen=True) @@ -58,12 +64,17 @@ def __init__( *, results_bucket: str | None = None, storage_adapter: StorageAdapter | None = None, + upload_concurrency: int | None = None, ) -> None: self._job_file_storage = JobFileStorage( storage_adapter=storage_adapter, results_bucket=results_bucket, ) self.results_bucket = self._job_file_storage.results_bucket + configured_concurrency = getattr(settings, "RESULT_UPLOAD_CONCURRENCY", 20) + self.upload_concurrency = self._validate_upload_concurrency( + configured_concurrency if upload_concurrency is None else upload_concurrency + ) def build_zip_key(self, *, job_id: str) -> str: return self._job_file_storage.build_result_zip_key(job_id=job_id) @@ -109,19 +120,33 @@ def upload( ) self._cleanup_file(zip_path) - raw_files: dict[str, str] = {} artifact_ref_filter = self._normalize_artifact_refs(artifact_refs) - for file_path in self._iter_raw_files(result_path): - relative_path = file_path.relative_to(result_path).as_posix() - if artifact_ref_filter is not None and relative_path not in artifact_ref_filter: - continue - raw_key = self.build_raw_key(job_id=job_id, relative_path=relative_path) - self._job_file_storage.upload_local_file( - str(file_path), - raw_key, - bucket=self.results_bucket, + upload_items = [ + ( + file_path, + file_path.relative_to(result_path).as_posix(), + ) + for file_path in self._iter_raw_files(result_path) + if artifact_ref_filter is None + or file_path.relative_to(result_path).as_posix() in artifact_ref_filter + ] + executor = self._get_upload_executor() + futures = [ + executor.submit( + self._upload_raw_file, + job_id=job_id, + file_path=file_path, + relative_path=relative_path, ) - raw_files[relative_path] = raw_key + for file_path, relative_path in upload_items + ] + try: + uploaded_items = [future.result() for future in futures] + except Exception: + for future in futures: + future.cancel() + raise + raw_files = dict(uploaded_items) return UploadedResultBundle( zip_key=zip_key, @@ -129,6 +154,40 @@ def upload( raw_files=raw_files, ) + def _upload_raw_file( + self, + *, + job_id: str, + file_path: Path, + relative_path: str, + ) -> tuple[str, str]: + raw_key = self.build_raw_key(job_id=job_id, relative_path=relative_path) + self._job_file_storage.upload_local_file( + str(file_path), + raw_key, + bucket=self.results_bucket, + ) + return relative_path, raw_key + + def _validate_upload_concurrency(self, upload_concurrency: int) -> int: + if not 1 <= upload_concurrency <= _MAX_RESULT_UPLOAD_CONCURRENCY: + raise ValueError( + "upload_concurrency must be between 1 and " + f"{_MAX_RESULT_UPLOAD_CONCURRENCY}" + ) + return upload_concurrency + + def _get_upload_executor(self) -> ThreadPoolExecutor: + with _upload_executor_lock: + executor = _upload_executors.get(self.upload_concurrency) + if executor is None: + executor = ThreadPoolExecutor( + max_workers=self.upload_concurrency, + thread_name_prefix="result-upload", + ) + _upload_executors[self.upload_concurrency] = executor + return executor + def generate_url(self, *, storage_key: str, expires_in: int = 3600) -> str | None: return self._job_file_storage.generate_download_url( storage_key, diff --git a/packages/shared-python/shared/tests/test_storage_config_contract.py b/packages/shared-python/shared/tests/test_storage_config_contract.py index d7dce29ef..8931688f0 100644 --- a/packages/shared-python/shared/tests/test_storage_config_contract.py +++ b/packages/shared-python/shared/tests/test_storage_config_contract.py @@ -43,6 +43,13 @@ def test_aws_s3_uses_default_credential_chain_when_keys_are_empty( assert client_arguments["region_name"] == "us-east-1" assert "aws_access_key_id" not in client_arguments assert "aws_secret_access_key" not in client_arguments + assert client_arguments["config"].max_pool_connections == 20 + + +def test_result_upload_concurrency_defaults_to_twenty() -> None: + config: StorageConfig = create_storage_config() + + assert config.RESULT_UPLOAD_CONCURRENCY == 20 def test_aws_s3_passes_complete_explicit_credentials(