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

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

The uvicorn supervisor imports the app before spawning workers, so the
module-level engine build left it holding an idle ninth copy of every model.
Engines now build in the FastAPI lifespan, which runs only in serving
processes. With more than one worker, each now exits after a jittered request
count and the supervisor respawns it, bounding per-worker memory growth that
otherwise reached the task limit about 11 days after a deploy.

* improvement(pii): move the start command into pii-entrypoint.sh and build engines in lifespan

Author: waleedlatif1

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"]