fix(pii): bound NER memory per worker so large texts stop OOM-killing Presidio workers

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,23 +332,24 @@ 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
         )
     # Regex-only requests reuse a blank NlpArtifacts to skip the spaCy NLP pass;
     # 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"}
 
 

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.