59 lines
2.0 KiB
Python
59 lines
2.0 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 logging
|
|
import multiprocessing
|
|
import os
|
|
from concurrent.futures import ProcessPoolExecutor
|
|
from threading import Lock
|
|
|
|
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(f"ebook_phrase_extraction_pool_started {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")
|