From 02c7960a8e14b376fabb0ba57c6626d88a3596f8 Mon Sep 17 00:00:00 2001 From: suguanYang Date: Sat, 12 Sep 2026 00:53:38 +0800 Subject: [PATCH 1/2] perf: bulk materialization publication writes --- .../app/services/demo/source_materializer.py | 122 ++++++++++----- .../shared/core/config/database.py | 16 ++ .../services/jobs/lifecycle/publication.py | 42 ++++-- .../shared/services/redis/__init__.py | 4 + .../services/redis/publication_semaphore.py | 142 ++++++++++++++++++ .../shared/services/redis/redis_service.py | 26 ++++ .../services/retrieval/map_unit_index.py | 71 ++++++--- 7 files changed, 351 insertions(+), 72 deletions(-) create mode 100644 packages/shared-python/shared/services/redis/publication_semaphore.py diff --git a/apps/api/app/services/demo/source_materializer.py b/apps/api/app/services/demo/source_materializer.py index 0f0bef48f..1e903f53d 100644 --- a/apps/api/app/services/demo/source_materializer.py +++ b/apps/api/app/services/demo/source_materializer.py @@ -4,6 +4,7 @@ import shutil import tempfile +import time from collections.abc import Iterable from dataclasses import dataclass from datetime import datetime, timedelta, timezone @@ -11,6 +12,8 @@ from pathlib import Path from uuid import uuid4 +import logfire + from app.services.demo.source_catalog import DemoSourceCatalog, DemoSourceDefinition from sqlalchemy.exc import IntegrityError from sqlalchemy import delete, func, select @@ -23,6 +26,8 @@ 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.redis import RedisPublicationSemaphore, RedisServiceFactory +from shared.core.config import settings from shared.services.storage.result_storage import get_result_storage @@ -50,6 +55,7 @@ def __init__( ) -> None: self._catalog = catalog self._publication_service = publication_service or RetrievalPublicationService() + self._redis_service = RedisServiceFactory.get_service() async def materialize_sources( self, @@ -93,12 +99,19 @@ async def materialize_sources( results: list[MaterializedDemoSource] = [] for source in selected_sources: try: + semaphore = RedisPublicationSemaphore( + self._redis_service, + concurrency=settings.MATERIALIZATION_DB_PUBLICATION_CONCURRENCY, + lease_seconds=settings.MATERIALIZATION_DB_PUBLICATION_LEASE_SECONDS, + acquire_timeout_seconds=settings.MATERIALIZATION_DB_PUBLICATION_ACQUIRE_TIMEOUT_SECONDS, + ) result = await self._materialize_source( db, user_id=user_id, namespace=namespace, source=source, claim=claims[source.demo_source_id], + semaphore=semaphore, ) results.append(result) except Exception: @@ -122,15 +135,22 @@ async def _materialize_source( namespace: str, source: DemoSourceDefinition, claim: DemoMaterialization, + semaphore: RedisPublicationSemaphore, ) -> MaterializedDemoSource: document_id = f"doc_{uuid4().hex[:12]}" job_id = f"job_demo_{uuid4().hex[:12]}" job_result_id = str(uuid4()) timestamp = _utc_now() + stage_started_at = time.perf_counter() result_bundle = _upload_demo_result_bundle( job_id=job_id, source_directory=self._catalog.source_directory(source), ) + logfire.info( + "Demo materialization source bundle upload completed", + demo_source_id=source.demo_source_id, + duration_seconds=time.perf_counter() - stage_started_at, + ) db.add( Job( @@ -170,45 +190,79 @@ async def _materialize_source( updated_at=timestamp, ) ) - await db.flush() chunks = self._catalog.publication_chunks(source) - 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, - ) + stage_started_at = time.perf_counter() + wait_seconds = await semaphore.acquire() + logfire.info( + "Demo materialization publication semaphore acquired", + demo_source_id=source.demo_source_id, + wait_seconds=wait_seconds, ) - await db.run_sync( - lambda sync_db: self._publication_service.publish_document_graph( - sync_db, - job_id=job_id, - job_result_id=job_result_id, + try: + base_rows_started_at = time.perf_counter() + await db.flush() + logfire.info( + "Demo materialization base rows completed", + demo_source_id=source.demo_source_id, + duration_seconds=time.perf_counter() - base_rows_started_at, ) - ) - await db.flush() - - 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, - document_id=document_id, + 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, - source_file_name=source.title, - ), - manifest_payload=manifest_payload, + chunks=[dict(chunk) for chunk in chunks], + update_namespace_snapshot=False, + ) ) - ) - await db.commit() + await db.run_sync( + lambda sync_db: self._publication_service.publish_document_graph( + sync_db, + job_id=job_id, + job_result_id=job_result_id, + ) + ) + await db.flush() + logfire.info( + "Demo materialization sections chunks and map index completed", + demo_source_id=source.demo_source_id, + duration_seconds=time.perf_counter() - stage_started_at, + chunk_count=len(chunks), + ) + + 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 + stage_started_at = time.perf_counter() + await db.run_sync( + lambda sync_db: self._publication_service.update_namespace_snapshot( + sync_db, + scope=DocumentPublicationScope( + user_id=user_id, + namespace=namespace, + document_id=document_id, + job_result_id=job_result_id, + source_file_name=source.title, + ), + manifest_payload=manifest_payload, + ) + ) + logfire.info( + "Demo materialization namespace snapshot completed", + demo_source_id=source.demo_source_id, + duration_seconds=time.perf_counter() - stage_started_at, + ) + commit_started_at = time.perf_counter() + await db.commit() + logfire.info( + "Demo materialization publication commit completed", + demo_source_id=source.demo_source_id, + duration_seconds=time.perf_counter() - commit_started_at, + ) + finally: + await semaphore.release() claim.document_id = document_id claim.status = "ready" diff --git a/packages/shared-python/shared/core/config/database.py b/packages/shared-python/shared/core/config/database.py index a9baac5f5..53d1b97d1 100644 --- a/packages/shared-python/shared/core/config/database.py +++ b/packages/shared-python/shared/core/config/database.py @@ -47,6 +47,22 @@ class DatabaseConfig(BaseModel): default=50, description="Celery gevent worker concurrency" ) + MATERIALIZATION_DB_PUBLICATION_CONCURRENCY: int = Field( + default=2, + ge=1, + description="Maximum concurrent materialization database publications", + ) + MATERIALIZATION_DB_PUBLICATION_LEASE_SECONDS: int = Field( + default=900, + ge=60, + description="Redis lease duration for a materialization publication permit", + ) + MATERIALIZATION_DB_PUBLICATION_ACQUIRE_TIMEOUT_SECONDS: float = Field( + default=30.0, + ge=0.1, + description="Maximum time to wait for a materialization publication permit", + ) + def get_ssl_connect_args(self) -> dict: """Return SSL connect args for psycopg2.""" ssl_args = {"sslmode": self.DB_SSL_MODE} diff --git a/packages/shared-python/shared/services/jobs/lifecycle/publication.py b/packages/shared-python/shared/services/jobs/lifecycle/publication.py index a8a3e7473..4ed9df81a 100644 --- a/packages/shared-python/shared/services/jobs/lifecycle/publication.py +++ b/packages/shared-python/shared/services/jobs/lifecycle/publication.py @@ -1,5 +1,6 @@ from __future__ import annotations +from contextlib import nullcontext from dataclasses import dataclass from typing import Any @@ -11,6 +12,8 @@ from shared.models.schemas.job_metadata import JobMetadataHelper from shared.models.schemas.retrieval_namespace import normalize_retrieval_namespace from shared.services.redis.redis_sync_service import SyncRedisServiceFactory +from shared.services.redis.publication_semaphore import SyncRedisPublicationSemaphore +from shared.core.config import settings from shared.services.retrieval.publication_service import RetrievalPublicationService from shared.services.retrieval.publication_models import ( ExistingDocumentScope, @@ -53,25 +56,38 @@ def publish_result( section_summaries: dict[str, str] | None, document_top_summary: str | None = None, ) -> JobPublicationOutcome: - previous_document_scope = self._retrieval_publication.get_existing_document_scope( - db, - job_id=job_id, + semaphore = SyncRedisPublicationSemaphore( + SyncRedisServiceFactory.get_service(), + concurrency=settings.MATERIALIZATION_DB_PUBLICATION_CONCURRENCY, + lease_seconds=settings.MATERIALIZATION_DB_PUBLICATION_LEASE_SECONDS, + acquire_timeout_seconds=settings.MATERIALIZATION_DB_PUBLICATION_ACQUIRE_TIMEOUT_SECONDS, ) - published_document_state = self._retrieval_publication.publish_document_state( - db, - job_id=job_id, - job_result_id=job_result_id, - chunks=chunks, - section_summaries=section_summaries, + job_type = db.execute( + select(Job.job_type).where(Job.job_id == job_id) + ).scalar_one_or_none() + publication_context = ( + semaphore if job_type == "demo_materialization" else nullcontext() ) - if _should_publish_document_graph(published_document_state): - assert published_document_state is not None - self._retrieval_publication.publish_document_graph( + with publication_context: + previous_document_scope = self._retrieval_publication.get_existing_document_scope( + db, + job_id=job_id, + ) + published_document_state = self._retrieval_publication.publish_document_state( db, job_id=job_id, job_result_id=job_result_id, - top_summary=document_top_summary, + chunks=chunks, + section_summaries=section_summaries, ) + if _should_publish_document_graph(published_document_state): + assert published_document_state is not None + self._retrieval_publication.publish_document_graph( + db, + job_id=job_id, + job_result_id=job_result_id, + top_summary=document_top_summary, + ) cache_invalidation = self._build_cache_invalidation( db, diff --git a/packages/shared-python/shared/services/redis/__init__.py b/packages/shared-python/shared/services/redis/__init__.py index a0cb3f46e..f61986a6e 100644 --- a/packages/shared-python/shared/services/redis/__init__.py +++ b/packages/shared-python/shared/services/redis/__init__.py @@ -20,6 +20,8 @@ __all__ = [ "RedisService", "RedisServiceFactory", + "RedisPublicationSemaphore", + "SyncRedisPublicationSemaphore", "RedisMonitor", "RedisAlertManager", "RedisAlertNotifier", @@ -38,6 +40,8 @@ _EXPORT_MODULES: dict[str, str] = { "RedisService": "shared.services.redis.redis_service", "RedisServiceFactory": "shared.services.redis.redis_service_factory", + "RedisPublicationSemaphore": "shared.services.redis.publication_semaphore", + "SyncRedisPublicationSemaphore": "shared.services.redis.publication_semaphore", "RedisMonitor": "shared.services.redis.redis_monitor", "RedisAlertManager": "shared.services.redis.redis_alerts", "RedisAlertNotifier": "shared.services.redis.redis_alerts", diff --git a/packages/shared-python/shared/services/redis/publication_semaphore.py b/packages/shared-python/shared/services/redis/publication_semaphore.py new file mode 100644 index 000000000..742b6f73e --- /dev/null +++ b/packages/shared-python/shared/services/redis/publication_semaphore.py @@ -0,0 +1,142 @@ +"""Cross-instance concurrency control for database publication.""" + +from __future__ import annotations + +import asyncio +import time +import uuid +from typing import Any + +from loguru import logger + +from shared.services.redis.redis_service import RedisService +from shared.services.redis.redis_sync_service import SyncRedisService + +_RELEASE_SCRIPT = """ +if redis.call('get', KEYS[1]) == ARGV[1] then + return redis.call('del', KEYS[1]) +end +return 0 +""" + + +class RedisPublicationSemaphore: + """Lease-based Redis semaphore using one key per permit.""" + + def __init__( + self, + redis_service: RedisService, + *, + concurrency: int, + lease_seconds: int, + acquire_timeout_seconds: float, + poll_interval_seconds: float = 0.05, + ) -> None: + if concurrency < 1: + raise ValueError("concurrency must be at least 1") + self._redis = redis_service + self._concurrency = concurrency + self._lease_seconds = lease_seconds + self._acquire_timeout_seconds = acquire_timeout_seconds + self._poll_interval_seconds = poll_interval_seconds + self._owner = uuid.uuid4().hex + self._permit_key: str | None = None + + async def acquire(self) -> float: + """Acquire a permit and return the time spent waiting in seconds.""" + started_at = time.perf_counter() + deadline = started_at + self._acquire_timeout_seconds + while time.perf_counter() < deadline: + for permit_number in range(self._concurrency): + key = f"lock:materialization_publication:{permit_number}" + acquired = await self._redis.set_nx( + key, + self._owner, + ex=self._lease_seconds, + ) + if acquired: + self._permit_key = key + return time.perf_counter() - started_at + await asyncio.sleep(self._poll_interval_seconds) + raise TimeoutError("Timed out waiting for a materialization publication permit") + + async def release(self) -> bool: + """Release this semaphore's permit only when still its owner.""" + if self._permit_key is None: + return False + key = self._permit_key + self._permit_key = None + try: + result = await self._redis.eval( + _RELEASE_SCRIPT, + keys=[key], + args=[self._owner], + ) + return bool(result) + except Exception as error: + logger.warning(f"Failed to release publication permit: {error}") + return False + + async def __aenter__(self) -> "RedisPublicationSemaphore": + await self.acquire() + return self + + async def __aexit__(self, exc_type: Any, exc_value: Any, traceback: Any) -> None: + await self.release() + + +class SyncRedisPublicationSemaphore: + """Synchronous lease-based semaphore for gevent worker publication.""" + + def __init__( + self, + redis_service: SyncRedisService, + *, + concurrency: int, + lease_seconds: int, + acquire_timeout_seconds: float, + poll_interval_seconds: float = 0.05, + ) -> None: + if concurrency < 1: + raise ValueError("concurrency must be at least 1") + self._redis = redis_service + self._concurrency = concurrency + self._lease_seconds = lease_seconds + self._acquire_timeout_seconds = acquire_timeout_seconds + self._poll_interval_seconds = poll_interval_seconds + self._owner = uuid.uuid4().hex + self._permit_key: str | None = None + + def acquire(self) -> float: + """Acquire a permit and return the time spent waiting in seconds.""" + started_at = time.perf_counter() + deadline = started_at + self._acquire_timeout_seconds + while time.perf_counter() < deadline: + for permit_number in range(self._concurrency): + key = f"lock:materialization_publication:{permit_number}" + if self._redis.set_nx(key, self._owner, ex=self._lease_seconds): + self._permit_key = key + return time.perf_counter() - started_at + time.sleep(self._poll_interval_seconds) + raise TimeoutError("Timed out waiting for a materialization publication permit") + + def release(self) -> bool: + """Release this semaphore's permit only when still its owner.""" + if self._permit_key is None: + return False + key = self._permit_key + self._permit_key = None + try: + return bool( + self._redis.eval(_RELEASE_SCRIPT, keys=[key], args=[self._owner]) + ) + except Exception as error: + logger.warning(f"Failed to release publication permit: {error}") + return False + + def __enter__(self) -> "SyncRedisPublicationSemaphore": + self.acquire() + return self + + def __exit__(self, exc_type: Any, exc_value: Any, traceback: Any) -> None: + self.release() diff --git a/packages/shared-python/shared/services/redis/redis_service.py b/packages/shared-python/shared/services/redis/redis_service.py index c25a3e1d6..cb45d75d4 100644 --- a/packages/shared-python/shared/services/redis/redis_service.py +++ b/packages/shared-python/shared/services/redis/redis_service.py @@ -298,6 +298,32 @@ async def _operation(): original_exception=e, ) + async def eval( + self, script: str, keys: list[str], args: list[Any] | None = None + ) -> Any: + """Execute a Lua script with consistently namespaced keys.""" + try: + client = await self._get_client() + full_keys = [self._build_key(key) for key in keys] + + async def _operation() -> Any: + return await _await_redis_result( + client.eval( + script, + len(full_keys), + *(full_keys + (args or [])), + ) + ) + + return await self._execute_with_retry(_operation) + except Exception as e: + logger.error(f"Redis EVAL operation failed: {e}") + raise RedisOperationError( + internal_message=f"EVAL operation failed: {str(e)}", + operation="EVAL", + original_exception=e, + ) + # ==================== List Operations ==================== async def lpush(self, key: str, *values: Any) -> int: diff --git a/packages/shared-python/shared/services/retrieval/map_unit_index.py b/packages/shared-python/shared/services/retrieval/map_unit_index.py index c58e5f8ef..19f95200e 100644 --- a/packages/shared-python/shared/services/retrieval/map_unit_index.py +++ b/packages/shared-python/shared/services/retrieval/map_unit_index.py @@ -6,7 +6,7 @@ from hashlib import sha256 from uuid import uuid4 -from sqlalchemy import delete, select +from sqlalchemy import Table, delete, insert, select from sqlalchemy.orm import Session from shared.models.database.document import ( @@ -29,6 +29,8 @@ __all__ = ["MAP_UNIT_INDEX_FORMAT_VERSION", "replace_document_map_units"] +_BULK_INSERT_BATCH_SIZE = 5_000 + def replace_document_map_units( db: Session, @@ -86,6 +88,8 @@ def replace_document_map_units( ) persisted_count = 0 token_count = 0 + map_unit_rows: list[dict[str, object]] = [] + token_rows: list[dict[str, object]] = [] path_unit_df: Counter[str] = Counter() content_unit_df: Counter[str] = Counter() path_document_count: int = 0 @@ -113,39 +117,43 @@ def replace_document_map_units( # embeds them), so this is the same ownership the query-time scorer # sees, not a new computation. section_chunk_types = {u.chunk_type for u in provider.self_units(section_id)} - db.add( - DocumentMapUnit( - id=map_unit_id, - document_id=scope.document_id, - job_result_id=scope.job_result_id, - unit_id=unit_id, - section_id=section_id, - unit_kind=str(unit.get("kind") or "leaf"), - path_token_count=len(path_tokens), - content_token_count=len(content_tokens), - term_search_text_lower=str(unit.get("term_search_text") or "").lower(), - has_image="image" in section_chunk_types, - has_table="table" in section_chunk_types, - sort_order=sort_order, - ) + map_unit_rows.append( + { + "id": map_unit_id, + "document_id": scope.document_id, + "job_result_id": scope.job_result_id, + "unit_id": unit_id, + "section_id": section_id, + "unit_kind": str(unit.get("kind") or "leaf"), + "path_token_count": len(path_tokens), + "content_token_count": len(content_tokens), + "term_search_text_lower": str( + unit.get("term_search_text") or "" + ).lower(), + "has_image": "image" in section_chunk_types, + "has_table": "table" in section_chunk_types, + "sort_order": sort_order, + } ) for channel, frequencies in ( ("path", Counter(path_tokens)), ("content", Counter(content_tokens)), ): for token, frequency in frequencies.items(): - db.add( - DocumentMapUnitToken( - id=f"dmut_{uuid4().hex[:31]}", - map_unit_id=map_unit_id, - channel=channel, - token=token, - token_hash=sha256(token.encode("utf-8")).hexdigest(), - frequency=frequency, - ) + token_rows.append( + { + "id": f"dmut_{uuid4().hex[:31]}", + "map_unit_id": map_unit_id, + "channel": channel, + "token": token, + "token_hash": sha256(token.encode("utf-8")).hexdigest(), + "frequency": frequency, + } ) token_count += len(frequencies) persisted_count += 1 + _execute_bulk_insert(db, DocumentMapUnit, map_unit_rows) + _execute_bulk_insert(db, DocumentMapUnitToken, token_rows) db.add( DocumentMapUnitIndex( id=f"dmui_{uuid4().hex}", @@ -170,6 +178,19 @@ def replace_document_map_units( ) +def _execute_bulk_insert( + db: Session, + model: type[DocumentMapUnit] | type[DocumentMapUnitToken], + rows: list[dict[str, object]], +) -> None: + """Insert derived index rows in bounded Core batches.""" + if not rows: + return + table: Table = model.__table__ + for start in range(0, len(rows), _BULK_INSERT_BATCH_SIZE): + db.execute(insert(table), rows[start : start + _BULK_INSERT_BATCH_SIZE]) + + def _to_section_row(section: DocumentSection) -> SectionRow: return SectionRow( section_id=section.section_id, From 92b830ab2cafa108272a1b3f633185931d106e31 Mon Sep 17 00:00:00 2001 From: suguanYang Date: Sat, 12 Sep 2026 13:38:13 +0800 Subject: [PATCH 2/2] fix: make semaphore exports explicit --- packages/shared-python/shared/services/redis/__init__.py | 4 ++++ 1 file changed, 4 insertions(+) diff --git a/packages/shared-python/shared/services/redis/__init__.py b/packages/shared-python/shared/services/redis/__init__.py index f61986a6e..befe58f38 100644 --- a/packages/shared-python/shared/services/redis/__init__.py +++ b/packages/shared-python/shared/services/redis/__init__.py @@ -11,6 +11,10 @@ from .key_builder import RedisKeyBuilder, RedisKeyType, redis_key_builder from .redis_alerts import AlertRule, RedisAlertManager, RedisAlertNotifier from .redis_monitor import RedisMonitor + from .publication_semaphore import ( + RedisPublicationSemaphore, + SyncRedisPublicationSemaphore, + ) from .redis_service import RedisService from .redis_service_factory import RedisServiceFactory from .retry_policy import RedisHealthChecker, RedisRetry