Richie
2706c4417d
Convert the ebook-search web app to async end to end and add concurrency to the protected-phrase extraction and judging pipeline so large books no longer block the event loop or the UI. ORM / infra: - Add get_async_postgres_engine and factor shared URL/connect_args building into build_postgres_url (reused by the sync and async engine builders) - Add async FastAPI session helpers (get_async_db, AsyncDbSession) with expire_on_commit=False to avoid implicit IO under asyncio App: - Use AsyncEngine/AsyncSession throughout routes, search, ingest, embeddings, answer, rerank and LLM calls; convert handlers to async - Share a single httpx.AsyncClient in app state for LLM requests; size the connection pool for concurrent phrase-judging workers - Add judge_tasks: run per-book judging as tracked background tasks so a book already being judged isn't double-queued Protected phrases: - Add a process pool (pool.py) and worker-count config (extraction/judge book/phrase workers) to parallelize candidate generation and judging - Split admin actions into all/missing variants for generation and judging Config: - Add protected_phrase_extraction_workers, phrase_judge_book_workers, phrase_judge_phrase_workers
102 lines
3.6 KiB
Python
102 lines
3.6 KiB
Python
"""Process pool for offloading CPU-bound phrase extraction off the request thread.
|
|
|
|
Phrase extraction is pure-Python CPU work (n-gram sliding, YAKE), so running it inline in a
|
|
sync request handler serializes concurrent recalculations behind the GIL. Submitting it to a
|
|
``ProcessPoolExecutor`` lets concurrent extractions run in parallel across cores instead. A
|
|
``spawn`` context is used so workers do not inherit the parent's database engine, connections,
|
|
or server threads.
|
|
"""
|
|
|
|
from __future__ import annotations
|
|
|
|
import asyncio
|
|
import logging
|
|
import multiprocessing
|
|
import os
|
|
from concurrent.futures import ProcessPoolExecutor
|
|
from threading import Lock
|
|
from typing import TYPE_CHECKING
|
|
|
|
from python.ebook_search.protected_phrases.extraction import extract_phrase_candidates_for_book
|
|
|
|
if TYPE_CHECKING:
|
|
from collections.abc import Mapping, Sequence
|
|
|
|
from python.ebook_search.config import EbookSearchConfig
|
|
from python.ebook_search.protected_phrases.models import PhraseCandidate
|
|
|
|
logger = logging.getLogger(__name__)
|
|
|
|
|
|
class _ExtractionPool:
|
|
"""Lazily created process-wide extraction pool and the lock guarding it."""
|
|
|
|
def __init__(self) -> None:
|
|
self.lock = Lock()
|
|
self.pool: ProcessPoolExecutor | None = None
|
|
|
|
|
|
_extraction_pool = _ExtractionPool()
|
|
|
|
|
|
def get_extraction_pool(max_workers: int) -> ProcessPoolExecutor:
|
|
"""Return the shared extraction process pool, creating it on first use.
|
|
|
|
Args:
|
|
max_workers (int): Desired worker count; values below 1 fall back to the CPU count.
|
|
|
|
Returns:
|
|
ProcessPoolExecutor: The shared pool for phrase extraction.
|
|
"""
|
|
with _extraction_pool.lock:
|
|
if _extraction_pool.pool is None:
|
|
workers = max_workers if max_workers > 0 else (os.cpu_count() or 1)
|
|
_extraction_pool.pool = ProcessPoolExecutor(
|
|
max_workers=workers,
|
|
mp_context=multiprocessing.get_context("spawn"),
|
|
)
|
|
logger.info("ebook_phrase_extraction_pool_started workers=%s", workers)
|
|
return _extraction_pool.pool
|
|
|
|
|
|
def shutdown_extraction_pool() -> None:
|
|
"""Shut down the shared extraction pool if it was started."""
|
|
with _extraction_pool.lock:
|
|
if _extraction_pool.pool is not None:
|
|
_extraction_pool.pool.shutdown(wait=False, cancel_futures=True)
|
|
_extraction_pool.pool = None
|
|
logger.info("ebook_phrase_extraction_pool_shutdown")
|
|
|
|
|
|
async def extract_phrase_candidates_in_pool(
|
|
book_text: str,
|
|
chapters: Sequence[str],
|
|
config: EbookSearchConfig,
|
|
*,
|
|
metadata: Mapping[str, object] | None,
|
|
) -> list[PhraseCandidate]:
|
|
"""Run book phrase extraction in a worker process and await the result.
|
|
|
|
Only the CPU-bound extraction runs in the worker; the caller keeps all database work in the
|
|
request process. The spaCy pipeline is not supported here because it is not picklable, so
|
|
this always runs the non-spaCy extraction path.
|
|
|
|
Args:
|
|
book_text (str): Full book text used for extraction.
|
|
chapters (Sequence[str]): Chapter-like text blocks used for frequency counts.
|
|
config (EbookSearchConfig): Runtime phrase-tuning settings.
|
|
metadata (Mapping[str, object] | None): Optional book metadata used as a candidate source.
|
|
|
|
Returns:
|
|
list[PhraseCandidate]: Scored candidates sorted best-first and capped per book.
|
|
"""
|
|
pool = get_extraction_pool(config.protected_phrase_extraction_workers)
|
|
future = pool.submit(
|
|
extract_phrase_candidates_for_book,
|
|
book_text,
|
|
list(chapters),
|
|
config,
|
|
metadata=dict(metadata) if metadata is not None else None,
|
|
)
|
|
return await asyncio.wrap_future(future)
|