diff --git a/apps/pii/server.py b/apps/pii/server.py index 405d96e13bc..cf8b347720d 100644 --- a/apps/pii/server.py +++ b/apps/pii/server.py @@ -8,9 +8,10 @@ import json import logging +import threading import time from collections.abc import AsyncIterator -from contextlib import asynccontextmanager +from contextlib import asynccontextmanager, nullcontext from typing import Any import regex as regex_module @@ -184,7 +185,6 @@ def build_analyzer() -> AnalyzerEngine: anonymizer: AnonymizerEngine NER_ENTITIES: frozenset _BLANK_ARTIFACTS: dict -_NLP_MAX_LENGTH: dict # Every entity the spaCy NER recognizers can actually produce. A request touching # any of these must run spaCy; a request naming only non-NER (regex/checksum) @@ -228,19 +228,22 @@ def _regex_only(entities: list[str] | None, score_threshold: float | None) -> bo ) -# spaCy refuses (error E088) any text longer than its model's `max_length`, so a -# longer text runs NER in windows of this many characters, following spaCy's -# guidance to segment long documents (it sizes NER memory at roughly 1GB per -# 100,000 characters of one text). +# A text longer than this runs NER in windows of this many characters, following +# spaCy's guidance to segment long documents. One NER pass costs memory that grows +# with the text's length (about 3GB for 900,000 characters), enough to get a worker +# OOM-killed, so no pass sees more than one window. That also keeps every pass under +# spaCy's `max_length`, which it refuses with error E088. NER_WINDOW_CHARS = 100_000 # How much of the previous window each window repeats. Far longer than any entity # spaCy tags, so an entity cut by one window's edge lies whole inside the next. NER_WINDOW_OVERLAP = 1_000 - -def _exceeds_nlp_length(text: str, language: str) -> bool: - limit = _NLP_MAX_LENGTH.get(language) - return limit is not None and len(text) > limit +# One NER pass (a batch, a short text, or one window) at a time per worker. FastAPI +# runs these sync endpoints on a threadpool, so without it a worker runs one pass per +# request in flight and their memory adds up until it is OOM-killed. A long text +# takes it per window, so it interleaves with other requests. Regex-only requests +# never take it. +_NER_SLOT = threading.Lock() def _window_entities( @@ -265,7 +268,7 @@ def _analyze_long( return_decision_process: bool = False, ad_hoc_recognizers: list[PatternRecognizer] | None = None, ) -> list[RecognizerResult]: - """Analyze a text longer than spaCy's `max_length`. + """Analyze a text longer than NER_WINDOW_CHARS. Pattern, checksum, and custom recognizers have no length limit, so they run once over the whole text (against blank NLP artifacts, as the regex-only path @@ -291,13 +294,15 @@ def _analyze_long( return results step = NER_WINDOW_CHARS - NER_WINDOW_OVERLAP for start in range(0, len(text), step): - for result in analyzer.analyze( - text=text[start : start + NER_WINDOW_CHARS], - language=language, - entities=window_entities, - score_threshold=score_threshold, - return_decision_process=return_decision_process, - ): + with _NER_SLOT: + window_results = analyzer.analyze( + text=text[start : start + NER_WINDOW_CHARS], + language=language, + entities=window_entities, + score_threshold=score_threshold, + return_decision_process=return_decision_process, + ) + for result in window_results: results.append( RecognizerResult( result.entity_type, @@ -327,7 +332,7 @@ def _analyze_one( ad_hoc_recognizers: list[PatternRecognizer] | None = None, ): regex_only = _regex_only(entities, score_threshold) - if not regex_only and _exceeds_nlp_length(text, language): + if not regex_only and len(text) > NER_WINDOW_CHARS: return _analyze_long( text, language, entities, score_threshold, return_decision_process, ad_hoc_recognizers ) @@ -335,15 +340,16 @@ def _analyze_one( # otherwise analyze() computes artifacts (runs spaCy) as usual. Custom-pattern # recognizers are regex-based, so they run fine against the blank artifacts. nlp_artifacts = _BLANK_ARTIFACTS.get(language) if regex_only else None - return analyzer.analyze( - text=text, - language=language, - entities=entities or None, - score_threshold=score_threshold, - return_decision_process=return_decision_process, - nlp_artifacts=nlp_artifacts, - ad_hoc_recognizers=ad_hoc_recognizers or None, - ) + with nullcontext() if regex_only else _NER_SLOT: + return analyzer.analyze( + text=text, + language=language, + entities=entities or None, + score_threshold=score_threshold, + return_decision_process=return_decision_process, + nlp_artifacts=nlp_artifacts, + ad_hoc_recognizers=ad_hoc_recognizers or None, + ) # Presidio's BatchAnalyzerEngine defaults nlp.pipe to batch_size=1, so a 2000-text @@ -359,7 +365,7 @@ def _analyze_many( ad_hoc_recognizers: list[PatternRecognizer] | None = None, ): """Analyze many texts, skipping the spaCy pass for regex-only requests and - analyzing any text longer than spaCy's `max_length` on its own (see + analyzing any text longer than NER_WINDOW_CHARS on its own (see `_analyze_long`).""" if _regex_only(entities, score_threshold): blank = _BLANK_ARTIFACTS.get(language) @@ -374,19 +380,20 @@ def _analyze_many( ) for text in texts ] - long_indices = {i for i, text in enumerate(texts) if _exceeds_nlp_length(text, language)} + long_indices = {i for i, text in enumerate(texts) if len(text) > NER_WINDOW_CHARS} short_indices = [i for i in range(len(texts)) if i not in long_indices] results: list[list[RecognizerResult]] = [[] for _ in texts] - short_results = batch_analyzer.analyze_iterator( - texts=[texts[i] for i in short_indices], - language=language, - batch_size=NLP_BATCH_SIZE, - entities=entities or None, - score_threshold=score_threshold, - ad_hoc_recognizers=ad_hoc_recognizers or None, - ) - for i, per_text in zip(short_indices, short_results): - results[i] = list(per_text) + with _NER_SLOT: + short_results = batch_analyzer.analyze_iterator( + texts=[texts[i] for i in short_indices], + language=language, + batch_size=NLP_BATCH_SIZE, + entities=entities or None, + score_threshold=score_threshold, + ad_hoc_recognizers=ad_hoc_recognizers or None, + ) + for i, per_text in zip(short_indices, short_results): + results[i] = list(per_text) for i in sorted(long_indices): results[i] = _analyze_long( texts[i], language, entities, score_threshold, ad_hoc_recognizers=ad_hoc_recognizers @@ -396,7 +403,7 @@ def _analyze_many( @asynccontextmanager async def lifespan(_app: FastAPI) -> AsyncIterator[None]: - global analyzer, batch_analyzer, anonymizer, NER_ENTITIES, _BLANK_ARTIFACTS, _NLP_MAX_LENGTH + global analyzer, batch_analyzer, anonymizer, NER_ENTITIES, _BLANK_ARTIFACTS logger.info("building analyzer (spacy)") analyzer = build_analyzer() batch_analyzer = BatchAnalyzerEngine(analyzer_engine=analyzer) @@ -423,10 +430,6 @@ async def lifespan(_app: FastAPI) -> AsyncIterator[None]: language: analyzer.nlp_engine.process_text("", language) for language in SUPPORTED_LANGUAGES } - _NLP_MAX_LENGTH = { - language: nlp.max_length - for language, nlp in getattr(analyzer.nlp_engine, "nlp", {}).items() - } yield @@ -686,8 +689,10 @@ def run_anonymize( ) +# Async so it answers on the event loop: requests queued on _NER_SLOT hold +# threadpool threads, and the load balancer's health check must not wait behind them. @app.get("/health") -def health() -> dict[str, str]: +async def health() -> dict[str, str]: return {"status": "ok"} diff --git a/docker/pii.Dockerfile b/docker/pii.Dockerfile index de731c96854..f2bc1af993a 100644 --- a/docker/pii.Dockerfile +++ b/docker/pii.Dockerfile @@ -61,6 +61,11 @@ USER pii # 5001 avoids colliding with the app's 3000 in local/compose runs on one host. EXPOSE 5001 +# glibc gives each threadpool thread its own malloc arena and rarely returns their +# memory, so a worker's footprint ratchets up to one NER pass's peak per thread +# that ran one. Two arenas keep it near a single pass. +ENV MALLOC_ARENA_MAX=2 + # Per-pattern regex match timeout (Presidio's `regex`-module `finditer(timeout=...)`), # an interactive backstop against a catastrophic user-supplied custom regex. Presidio # logs and skips a pattern that exceeds it — the request stays up rather than hanging.