fix(pii): load models only in workers and recycle them before memory runs out

apps/pii/server.py

@@ -9,6 +9,8 @@
 import json
 import logging
 import time
+from collections.abc import AsyncIterator
+from contextlib import asynccontextmanager
 from typing import Any
 
 import regex as regex_module
@@ -174,10 +176,14 @@ def build_analyzer() -> AnalyzerEngine:
     logger.addHandler(_log_handler)
     logger.propagate = False
 
-logger.info("building analyzer (spacy)")
-analyzer = build_analyzer()
-batch_analyzer = BatchAnalyzerEngine(analyzer_engine=analyzer)
-anonymizer = AnonymizerEngine()
+# Built per worker by `lifespan`, never at import: with --workers N the uvicorn
+# supervisor imports this module too, and a module-level build left it holding an
+# idle copy of every model.
+analyzer: AnalyzerEngine
+batch_analyzer: BatchAnalyzerEngine
+anonymizer: AnonymizerEngine
+NER_ENTITIES: frozenset
+_BLANK_ARTIFACTS: 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)
@@ -207,30 +213,6 @@ def _producible_ner_entities() -> frozenset:
     return frozenset(producible)
 
 
-NER_ENTITIES = _SPACY_NER_FLOOR | (
-    frozenset(
-        entity
-        for recognizer in analyzer.registry.recognizers
-        if isinstance(recognizer, SpacyRecognizer)
-        for entity in recognizer.supported_entities
-    )
-    & _producible_ner_entities()
-)
-
-# One blank NlpArtifacts per language, built once at startup. Passing these to
-# analyze() skips nlp_engine.process_text (the spaCy tok2vec+ner pass) entirely:
-# the pattern recognizers still match on the raw text, SpacyRecognizer finds
-# nothing in blank artifacts (and can only be consulted at all for claimed-but-
-# unproducible entities like PHONE_NUMBER), and score_threshold is unset so
-# detection is identical.
-# Only context-based score boosting (which needs real tokens) is unavailable — an
-# accepted trade for skipping NER on the hot block-output path. Read-only, so it
-# is safe to share across requests and workers.
-_BLANK_ARTIFACTS = {
-    language: analyzer.nlp_engine.process_text("", language)
-    for language in SUPPORTED_LANGUAGES
-}
-
 
 def _regex_only(entities: list[str] | None, score_threshold: float | None) -> bool:
     """True when the spaCy NLP pass can be skipped: the request names entities, none
@@ -308,7 +290,39 @@ def _analyze_many(
     )
 
 
-app = FastAPI(title="Sim Presidio", docs_url=None, redoc_url=None)
+@asynccontextmanager
+async def lifespan(_app: FastAPI) -> AsyncIterator[None]:
+    global analyzer, batch_analyzer, anonymizer, NER_ENTITIES, _BLANK_ARTIFACTS
+    logger.info("building analyzer (spacy)")
+    analyzer = build_analyzer()
+    batch_analyzer = BatchAnalyzerEngine(analyzer_engine=analyzer)
+    anonymizer = AnonymizerEngine()
+    NER_ENTITIES = _SPACY_NER_FLOOR | (
+        frozenset(
+            entity
+            for recognizer in analyzer.registry.recognizers
+            if isinstance(recognizer, SpacyRecognizer)
+            for entity in recognizer.supported_entities
+        )
+        & _producible_ner_entities()
+    )
+    # One blank NlpArtifacts per language, built once at startup. Passing these to
+    # analyze() skips nlp_engine.process_text (the spaCy tok2vec+ner pass) entirely:
+    # the pattern recognizers still match on the raw text, SpacyRecognizer finds
+    # nothing in blank artifacts (and can only be consulted at all for claimed-but-
+    # unproducible entities like PHONE_NUMBER), and score_threshold is unset so
+    # detection is identical.
+    # Only context-based score boosting (which needs real tokens) is unavailable — an
+    # accepted trade for skipping NER on the hot block-output path. Read-only, so it
+    # is safe to share across requests.
+    _BLANK_ARTIFACTS = {
+        language: analyzer.nlp_engine.process_text("", language)
+        for language in SUPPORTED_LANGUAGES
+    }
+    yield
+
+
+app = FastAPI(title="Sim Presidio", docs_url=None, redoc_url=None, lifespan=lifespan)
 
 
 @app.exception_handler(RequestValidationError)

docker/pii-entrypoint.sh

@@ -0,0 +1,14 @@
+#!/bin/sh
+# Starts uvicorn as PID 1 (exec) so SIGTERM reaches it directly. Quoted expansions
+# make a malformed PII_WORKERS fail uvicorn's arg parsing, not the shell.
+set -- --host 0.0.0.0 --port 5001 --workers "${PII_WORKERS:-1}" --timeout-graceful-shutdown 300
+
+# Worker memory grows with uptime, so with a supervisor (more than one worker) each
+# worker exits after a jittered request count, health checks included, and is
+# respawned. A single worker has no supervisor to respawn it, so it never recycles.
+if [ "${PII_WORKERS:-1}" -gt 1 ] 2>/dev/null; then
+  set -- "$@" --limit-max-requests "${PII_MAX_REQUESTS:-4000}" \
+    --limit-max-requests-jitter "${PII_MAX_REQUESTS_JITTER:-2000}"
+fi
+
+exec uvicorn server:app "$@"

docker/pii.Dockerfile

@@ -39,6 +39,7 @@ RUN groupadd -g 1001 pii && \
     chown -R pii:pii /app
 
 COPY --chown=pii:pii apps/pii/server.py ./
+COPY --chown=pii:pii --chmod=755 docker/pii-entrypoint.sh ./
 
 USER pii
 
@@ -63,7 +64,4 @@ HEALTHCHECK --interval=30s --timeout=5s --start-period=300s --retries=3 \
 # Worker count is env-driven so ONE image scales per task size: set PII_WORKERS to
 # the task's vCPU count (each worker loads the models independently, ~3 GB each, so
 # size task memory ≈ PII_WORKERS × 3 GB + overhead). Defaults to 1 for local/small.
-# `sh -c exec` expands the env var while keeping uvicorn as PID 1 for clean SIGTERM.
-# Quote the expansion so a malformed PII_WORKERS fails uvicorn arg-parsing rather
-# than being interpreted by the shell.
-CMD ["sh", "-c", "exec uvicorn server:app --host 0.0.0.0 --port 5001 --workers \"${PII_WORKERS:-1}\""]
+CMD ["./pii-entrypoint.sh"]