From 4626999f8e9e5104319444b098831317258dbe9b Mon Sep 17 00:00:00 2001 From: Theodore Li Date: Tue, 7 Jul 2026 16:23:52 -0700 Subject: [PATCH 1/2] feat(pii): add opt-in GLiNER NER engine (PII_ENGINE), device-agnostic MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit Swap the 4 NER entity types (PERSON/LOCATION/NRP/DATE_TIME) to a single multilingual GLiNER zero-shot model when PII_ENGINE=gliner; spaCy stays the default and all ~36 regex/checksum recognizers are identical on both engines. Device-agnostic via PII_DEVICE / cuda auto-detect — same code on Fargate CPU now and EC2-GPU later. - engines.py: side-effect-free builders; SharedModelGLiNERRecognizer loads ONE model shared across the 5 per-language instances and restricts labels to the entities it owns; small spaCy models keep tokenization/lemmas for the regex recognizers; fail-fast on the lean image - pii.Dockerfile: multi-stage — default target unchanged (lean spaCy); --target gliner is a superset (torch CPU + gliner + baked model) where both engines work; gliner-gpu scaffold for the GPU fleet - CI publishes the gliner variant (:staging-gliner/:latest-gliner, amd64) - Helm: pii.engine / pii.device values wired to PII_ENGINE/PII_DEVICE - scripts/bench_engines.py: throughput + NER-parity diff harness - tests: unit (mocked GLiNER) + in-image integration for both engines Co-Authored-By: Claude Fable 5 Claude-Session: https://claude.ai/code/session_01Up3F97mjCH9HCj1pX4J8VJ --- .github/workflows/ci.yml | 38 +++- .gitignore | 4 + apps/pii/engines.py | 266 +++++++++++++++++++++++++ apps/pii/requirements-dev.txt | 5 + apps/pii/requirements-gliner.txt | 10 + apps/pii/scripts/bench_engines.py | 206 +++++++++++++++++++ apps/pii/scripts/bench_payload.json | 97 +++++++++ apps/pii/server.py | 146 +++----------- apps/pii/tests/conftest.py | 5 + apps/pii/tests/test_engines.py | 105 ++++++++++ apps/pii/tests/test_integration.py | 102 ++++++++++ docker/pii.Dockerfile | 125 ++++++++++-- helm/sim/templates/deployment-pii.yaml | 10 +- helm/sim/tests/pii_test.yaml | 59 ++++++ helm/sim/values.schema.json | 9 + helm/sim/values.yaml | 13 ++ 16 files changed, 1056 insertions(+), 144 deletions(-) create mode 100644 apps/pii/engines.py create mode 100644 apps/pii/requirements-dev.txt create mode 100644 apps/pii/requirements-gliner.txt create mode 100644 apps/pii/scripts/bench_engines.py create mode 100644 apps/pii/scripts/bench_payload.json create mode 100644 apps/pii/tests/conftest.py create mode 100644 apps/pii/tests/test_engines.py create mode 100644 apps/pii/tests/test_integration.py diff --git a/.github/workflows/ci.yml b/.github/workflows/ci.yml index 25d8348d8c4..d2137803c05 100644 --- a/.github/workflows/ci.yml +++ b/.github/workflows/ci.yml @@ -90,6 +90,12 @@ jobs: ecr_repo_secret: ECR_REALTIME - dockerfile: ./docker/pii.Dockerfile ecr_repo_secret: ECR_PII + # Opt-in GLiNER variant of the pii image (torch + gliner + baked + # model). Deployed by setting PII_ENGINE=gliner on the pii service. + - dockerfile: ./docker/pii.Dockerfile + ecr_repo_secret: ECR_PII + target: gliner + tag_suffix: -gliner steps: - name: Checkout code uses: actions/checkout@df4cb1c069e1874edd31b4311f1884172cec0e10 # v6 @@ -126,7 +132,8 @@ jobs: file: ${{ matrix.dockerfile }} platforms: linux/amd64 push: true - tags: ${{ steps.login-ecr.outputs.registry }}/${{ steps.ecr-repo.outputs.name }}:dev + tags: ${{ steps.login-ecr.outputs.registry }}/${{ steps.ecr-repo.outputs.name }}:dev${{ matrix.tag_suffix || '' }} + target: ${{ matrix.target || '' }} provenance: false sbom: false @@ -201,6 +208,13 @@ jobs: - dockerfile: ./docker/pii.Dockerfile ghcr_image: ghcr.io/simstudioai/pii ecr_repo_secret: ECR_PII + # Opt-in GLiNER variant of the pii image (torch + gliner + baked + # model). amd64 only — the deploy fleet is amd64; no arm64/manifest. + - dockerfile: ./docker/pii.Dockerfile + ghcr_image: ghcr.io/simstudioai/pii + ecr_repo_secret: ECR_PII + target: gliner + tag_suffix: -gliner steps: - name: Checkout code uses: actions/checkout@df4cb1c069e1874edd31b4311f1884172cec0e10 # v6 @@ -244,26 +258,35 @@ jobs: ECR_REGISTRY="${{ steps.login-ecr.outputs.registry }}" ECR_REPO="${{ steps.ecr-repo.outputs.name }}" GHCR_IMAGE="${{ matrix.ghcr_image }}" + TAG_SUFFIX="${{ matrix.tag_suffix }}" + + # Suffixed variants (e.g. pii -gliner) are amd64-only and skip the + # manifest merge, so their GHCR tags are plain — no -amd64 alias. + if [ -n "$TAG_SUFFIX" ]; then + ARCH_SUFFIX="" + else + ARCH_SUFFIX="-amd64" + fi if [ "${{ github.ref }}" = "refs/heads/main" ]; then - ECR_TAG="latest" + ECR_TAG="latest${TAG_SUFFIX}" else - ECR_TAG="staging" + ECR_TAG="staging${TAG_SUFFIX}" fi ECR_IMAGE="${ECR_REGISTRY}/${ECR_REPO}:${ECR_TAG}" TAGS="${ECR_IMAGE}" if [ "${{ github.ref }}" = "refs/heads/main" ] && [ -n "$GHCR_IMAGE" ]; then - GHCR_AMD64="${GHCR_IMAGE}:latest-amd64" - GHCR_SHA="${GHCR_IMAGE}:${{ github.sha }}-amd64" + GHCR_AMD64="${GHCR_IMAGE}:latest${TAG_SUFFIX}${ARCH_SUFFIX}" + GHCR_SHA="${GHCR_IMAGE}:${{ github.sha }}${TAG_SUFFIX}${ARCH_SUFFIX}" TAGS="${TAGS},$GHCR_AMD64,$GHCR_SHA" if [ "${{ needs.detect-version.outputs.is_release }}" = "true" ]; then VERSION="${{ needs.detect-version.outputs.version }}" - GHCR_VERSION="${GHCR_IMAGE}:${VERSION}-amd64" + GHCR_VERSION="${GHCR_IMAGE}:${VERSION}${TAG_SUFFIX}${ARCH_SUFFIX}" TAGS="${TAGS},$GHCR_VERSION" - echo "📦 Adding version tag: ${VERSION}-amd64" + echo "📦 Adding version tag: ${VERSION}${TAG_SUFFIX}${ARCH_SUFFIX}" fi fi @@ -277,6 +300,7 @@ jobs: platforms: linux/amd64 push: true tags: ${{ steps.meta.outputs.tags }} + target: ${{ matrix.target || '' }} provenance: false sbom: false diff --git a/.gitignore b/.gitignore index a700a66602a..a8a0d8e2fde 100644 --- a/.gitignore +++ b/.gitignore @@ -96,3 +96,7 @@ i18n.cache # Personal Cursor Skills .cursor/skills/ask-sim/ + +# Python (apps/pii tests/tooling) +__pycache__/ +.pytest_cache/ diff --git a/apps/pii/engines.py b/apps/pii/engines.py new file mode 100644 index 00000000000..8af7184d26c --- /dev/null +++ b/apps/pii/engines.py @@ -0,0 +1,266 @@ +"""Analyzer engine builders for the PII service. + +Two NER engines share one recognizer surface: + +- spacy (default): the 5 large spaCy models do NER (PERSON/LOCATION/NRP/ + DATE_TIME) and tokenization. +- gliner (opt-in): one multilingual GLiNER model does NER on CPU or GPU; + small spaCy models remain only for tokenization + lemmas. + +Both engines register the identical regex/checksum recognizer set (Presidio +defaults, EXTRA_RECOGNIZERS, VIN) — only the source of the 4 NER entity types +differs. Side-effect free: importing this module loads no models. +""" + +import importlib.util + +import spacy.util +from presidio_analyzer import AnalyzerEngine, Pattern, PatternRecognizer +from presidio_analyzer.nlp_engine import NlpEngineProvider +from presidio_analyzer.predefined_recognizers import ( + AuAbnRecognizer, + AuAcnRecognizer, + AuMedicareRecognizer, + AuTfnRecognizer, + EsNieRecognizer, + EsNifRecognizer, + FiPersonalIdentityCodeRecognizer, + GLiNERRecognizer, + InAadhaarRecognizer, + InPanRecognizer, + InPassportRecognizer, + InVehicleRegistrationRecognizer, + InVoterRecognizer, + ItDriverLicenseRecognizer, + ItFiscalCodeRecognizer, + ItIdentityCardRecognizer, + ItPassportRecognizer, + ItVatCodeRecognizer, + PlPeselRecognizer, + SgFinRecognizer, + SgUenRecognizer, + UkNinoRecognizer, +) + +# Languages served. Each needs its spaCy model installed in the image; the +# es/it/pl/fi predefined recognizers (ES_NIF, IT_FISCAL_CODE, PL_PESEL, ...) +# auto-load once their NLP engine is present. +NLP_CONFIGURATION = { + "nlp_engine_name": "spacy", + "models": [ + {"lang_code": "en", "model_name": "en_core_web_lg"}, + {"lang_code": "es", "model_name": "es_core_news_lg"}, + {"lang_code": "it", "model_name": "it_core_news_lg"}, + {"lang_code": "pl", "model_name": "pl_core_news_lg"}, + {"lang_code": "fi", "model_name": "fi_core_news_lg"}, + ], +} +SUPPORTED_LANGUAGES = [m["lang_code"] for m in NLP_CONFIGURATION["models"]] + +# The gliner engine still needs a spaCy pipeline per language: the regex +# recognizers consume NlpArtifacts and the LemmaContextAwareEnhancer boosts +# scores from surrounding lemmas. The small models (~12-40MB each vs ~400MB +# large) keep tokenization + lemmas intact while GLiNER owns NER. Blank +# pipelines ("blank:xx") are not an option: Presidio's SpacyNlpEngine treats +# unknown model names as pip packages and tries to download them. +# labels_to_ignore strips the small models' NER output from NlpArtifacts — +# correctness comes from removing SpacyRecognizer in build_gliner_analyzer; +# this only silences unmapped-label noise. +GLINER_NLP_CONFIGURATION = { + "nlp_engine_name": "spacy", + "models": [ + {"lang_code": "en", "model_name": "en_core_web_sm"}, + {"lang_code": "es", "model_name": "es_core_news_sm"}, + {"lang_code": "it", "model_name": "it_core_news_sm"}, + {"lang_code": "pl", "model_name": "pl_core_news_sm"}, + {"lang_code": "fi", "model_name": "fi_core_news_sm"}, + ], + "ner_model_configuration": { + "labels_to_ignore": [ + "CARDINAL", "DATE", "EVENT", "FAC", "GPE", "LANGUAGE", "LAW", + "LOC", "MISC", "MONEY", "NORP", "ORDINAL", "ORG", "PER", + "PERCENT", "PERSON", "PRODUCT", "QUANTITY", "TIME", "WORK_OF_ART", + ], + }, +} + +# Zero-shot label prompts -> the 4 Presidio NER entities GLiNER owns. Multiple +# prompts per entity trade a little inference cost for recall; tune against +# scripts/bench_engines.py output. +GLINER_ENTITY_MAPPING = { + "person": "PERSON", + "name": "PERSON", + "location": "LOCATION", + "address": "LOCATION", + "date": "DATE_TIME", + "time": "DATE_TIME", + "nationality": "NRP", + "religious group": "NRP", + "political group": "NRP", + "ethnic group": "NRP", +} + +# Predefined recognizers Presidio ships but does NOT load into the default +# registry — they must be added explicitly. Each carries its own +# supported_language, so it fires under that language once its NLP model is +# loaded. en: UK/AU/IN/SG locale ids; es/it/pl/fi: national ids. +EXTRA_RECOGNIZERS = [ + UkNinoRecognizer, + AuAbnRecognizer, + AuAcnRecognizer, + AuTfnRecognizer, + AuMedicareRecognizer, + InPanRecognizer, + InAadhaarRecognizer, + InVehicleRegistrationRecognizer, + InVoterRecognizer, + InPassportRecognizer, + SgFinRecognizer, + SgUenRecognizer, + EsNifRecognizer, + EsNieRecognizer, + ItFiscalCodeRecognizer, + ItDriverLicenseRecognizer, + ItVatCodeRecognizer, + ItPassportRecognizer, + ItIdentityCardRecognizer, + PlPeselRecognizer, + FiPersonalIdentityCodeRecognizer, +] + + +class VinRecognizer(PatternRecognizer): + """VIN (17 chars, A-Z/0-9 excluding I/O/Q) with ISO 3779 check-digit + validation (position 9). Validation makes accidental matches on arbitrary + 17-char codes (request ids, SKUs, tokens) extremely unlikely. Some + non-North-American VINs omit the check digit and are skipped — an + intentional bias toward precision. + """ + + _TRANSLIT = { + **{str(d): d for d in range(10)}, + "A": 1, "B": 2, "C": 3, "D": 4, "E": 5, "F": 6, "G": 7, "H": 8, + "J": 1, "K": 2, "L": 3, "M": 4, "N": 5, "P": 7, "R": 9, + "S": 2, "T": 3, "U": 4, "V": 5, "W": 6, "X": 7, "Y": 8, "Z": 9, + } + _WEIGHTS = [8, 7, 6, 5, 4, 3, 2, 10, 0, 9, 8, 7, 6, 5, 4, 3, 2] + + def validate_result(self, pattern_text: str): + vin = pattern_text.upper() + if len(vin) != 17: + return False + try: + total = sum(self._TRANSLIT[c] * w for c, w in zip(vin, self._WEIGHTS)) + except KeyError: + return False + check = total % 11 + expected = "X" if check == 10 else str(check) + return vin[8] == expected + + +class SharedModelGLiNERRecognizer(GLiNERRecognizer): + """Per-language GLiNER recognizer sharing ONE loaded model. + + Presidio routes recognizers by supported_language, so the registry holds + one instance per served language — but each instance's load() would pull + its own ~1.2GB model copy. The first instance loads (an ImportError from + a missing gliner package propagates — fail fast in the lean image); the + rest reuse the cached model. + """ + + _shared_models: dict = {} + + def load(self) -> None: + key = (self.model_name, self.map_location) + cached = self._shared_models.get(key) + if cached is None: + super().load() + self._shared_models[key] = self.gliner + else: + self.gliner = cached + + def analyze(self, text, entities, nlp_artifacts=None): + """GLiNERRecognizer appends any requested entity it doesn't know as an + ad-hoc zero-shot label and returns its hits. The analyzer passes ALL + supported entities (~40) when a request doesn't narrow them, which + would prompt GLiNER for CREDIT_CARD/VIN/ES_NIF/... — wrong scope, and + inference cost scales with label count. Restrict to the NER entities + this recognizer owns.""" + requested = [e for e in (entities or self.supported_entities) if e in self.supported_entities] + if not requested: + return [] + return super().analyze(text, requested, nlp_artifacts) + + +def _register_common_recognizers(analyzer: AnalyzerEngine) -> None: + """Regex/checksum recognizers shared by both engines.""" + # VIN is language-agnostic, so register it under every served language — + # a recognizer only fires for the language the caller routes to. + vin_pattern = Pattern(name="vin", regex=r"\b[A-HJ-NPR-Z0-9]{17}\b", score=0.7) + for language in SUPPORTED_LANGUAGES: + analyzer.registry.add_recognizer( + VinRecognizer( + supported_entity="VIN", + patterns=[vin_pattern], + context=["vin", "vehicle", "chassis"], + supported_language=language, + ) + ) + for recognizer_cls in EXTRA_RECOGNIZERS: + analyzer.registry.add_recognizer(recognizer_cls()) + + +def build_spacy_analyzer() -> AnalyzerEngine: + nlp_engine = NlpEngineProvider(nlp_configuration=NLP_CONFIGURATION).create_engine() + analyzer = AnalyzerEngine(nlp_engine=nlp_engine, supported_languages=SUPPORTED_LANGUAGES) + _register_common_recognizers(analyzer) + return analyzer + + +def build_gliner_analyzer(model_name: str, device: str | None) -> AnalyzerEngine: + """GLiNER engine: one multilingual zero-shot model replaces spaCy NER for + PERSON/LOCATION/NRP/DATE_TIME; everything else is unchanged. + + :param model_name: HuggingFace id of the GLiNER model. + :param device: torch device ("cpu", "cuda", "cuda:0"); None auto-detects + via Presidio's device_detector (cuda when available, else cpu). + """ + # Fail fast with an actionable message on the lean image. Without these + # checks Presidio would try to pip-download the missing spaCy models at + # startup (a silent network fallback that dies with an unrelated pip + # permission error), and the gliner ImportError would surface only later. + if importlib.util.find_spec("gliner") is None: + raise RuntimeError( + "PII_ENGINE=gliner requires the gliner image variant " + "(docker build --target gliner); the gliner package is not installed" + ) + missing = [ + m["model_name"] + for m in GLINER_NLP_CONFIGURATION["models"] + if not spacy.util.is_package(m["model_name"]) + ] + if missing: + raise RuntimeError( + f"PII_ENGINE=gliner needs spaCy models {missing}; " + "use the gliner image variant (docker build --target gliner)" + ) + nlp_engine = NlpEngineProvider(nlp_configuration=GLINER_NLP_CONFIGURATION).create_engine() + analyzer = AnalyzerEngine(nlp_engine=nlp_engine, supported_languages=SUPPORTED_LANGUAGES) + # The default registry wires SpacyRecognizer per language; with GLiNER + # owning the NER entities it would emit duplicate/competing spans from the + # small models' ner pipe. remove_recognizer only logs when nothing matched, + # so assert the removal actually happened. + analyzer.registry.remove_recognizer("SpacyRecognizer") + if any(r.name == "SpacyRecognizer" for r in analyzer.registry.recognizers): + raise RuntimeError("SpacyRecognizer removal failed; Presidio registry layout changed") + for language in SUPPORTED_LANGUAGES: + analyzer.registry.add_recognizer( + SharedModelGLiNERRecognizer( + entity_mapping=GLINER_ENTITY_MAPPING, + model_name=model_name, + map_location=device, + supported_language=language, + ) + ) + _register_common_recognizers(analyzer) + return analyzer diff --git a/apps/pii/requirements-dev.txt b/apps/pii/requirements-dev.txt new file mode 100644 index 00000000000..4c330d999ae --- /dev/null +++ b/apps/pii/requirements-dev.txt @@ -0,0 +1,5 @@ +# Test-only deps. Unit tests need requirements.txt + this file (no models); +# integration tests additionally need the models baked into the docker images +# (see tests/test_integration.py). +pytest==8.4.1 +httpx==0.28.1 diff --git a/apps/pii/requirements-gliner.txt b/apps/pii/requirements-gliner.txt new file mode 100644 index 00000000000..d620c5f1825 --- /dev/null +++ b/apps/pii/requirements-gliner.txt @@ -0,0 +1,10 @@ +# Extras for the opt-in GLiNER engine — installed ONLY in the `gliner` +# Dockerfile target, on top of requirements.txt. Pinned for reproducible image +# builds; bump deliberately. presidio-analyzer 2.2.362 requires +# gliner >=0.2.13,<1.0.0. +# +# torch is pinned in the Dockerfile instead: the CPU and CUDA targets install +# the same version from different wheel indexes. +gliner==0.2.27 +transformers==4.56.2 +huggingface_hub==0.35.3 diff --git a/apps/pii/scripts/bench_engines.py b/apps/pii/scripts/bench_engines.py new file mode 100644 index 00000000000..8d902fe4d22 --- /dev/null +++ b/apps/pii/scripts/bench_engines.py @@ -0,0 +1,206 @@ +"""Benchmark + parity harness for the spacy vs gliner NER engines. + +Runs the same payload through both engines and reports per-engine throughput +(batch analyze, the production /redact_batch path) and per-text latency, plus +an accuracy diff over the 4 NER entity types (PERSON/LOCATION/NRP/DATE_TIME). +Non-NER (regex/checksum) results must be identical between engines — both +register the same recognizers — so any mismatch there is a wiring bug and the +script exits non-zero. + +Meant to run inside the gliner image (the only one with both engines): + + docker run --rm sim-pii:gliner python scripts/bench_engines.py + docker run --rm -v $PWD/texts.json:/data.json sim-pii:gliner \\ + python scripts/bench_engines.py --payload /data.json + +Payload format: JSON list of {"text": str, "language": str} objects. +This doubles as the tuning harness for GLINER_ENTITY_MAPPING label prompts. +""" + +import argparse +import json +import statistics +import sys +import time +from collections import defaultdict +from pathlib import Path + +sys.path.insert(0, str(Path(__file__).resolve().parent.parent)) + +import engines # noqa: E402 + +# Entities sourced from the NER models rather than regex/checksum patterns. +# ORGANIZATION is emitted by the spacy engine's NER on unfiltered requests but +# is not in the app's supported set and has no GLiNER mapping — it shows up in +# the NER diff (spacy-only) rather than failing the regex-parity gate. +NER_ENTITIES = {"PERSON", "LOCATION", "NRP", "DATE_TIME", "ORGANIZATION"} +DEFAULT_PAYLOAD = Path(__file__).resolve().parent / "bench_payload.json" + + +def parse_args(): + parser = argparse.ArgumentParser(description=__doc__.splitlines()[0]) + parser.add_argument("--payload", type=Path, default=DEFAULT_PAYLOAD) + parser.add_argument("--engines", default="spacy,gliner") + parser.add_argument("--runs", type=int, default=3) + parser.add_argument("--warmup", type=int, default=1) + parser.add_argument("--device", default=None, help="torch device for gliner (default: auto)") + parser.add_argument("--gliner-model", default="urchade/gliner_multi_pii-v1") + parser.add_argument("--max-examples", type=int, default=10) + parser.add_argument("--json", action="store_true", help="emit machine-readable JSON") + return parser.parse_args() + + +def build(engine: str, args) -> tuple: + started = time.perf_counter() + if engine == "spacy": + analyzer = engines.build_spacy_analyzer() + elif engine == "gliner": + analyzer = engines.build_gliner_analyzer(model_name=args.gliner_model, device=args.device) + else: + raise ValueError(f"Unknown engine {engine!r}") + return analyzer, time.perf_counter() - started + + +def analyze_all(analyzer, items) -> list[list]: + """One analyze() call per text, in payload order.""" + return [analyzer.analyze(text=item["text"], language=item["language"]) for item in items] + + +def bench(analyzer, items, runs: int, warmup: int) -> dict: + for _ in range(warmup): + analyze_all(analyzer, items) + run_times = [] + latencies = [] + for _ in range(runs): + run_started = time.perf_counter() + for item in items: + text_started = time.perf_counter() + analyzer.analyze(text=item["text"], language=item["language"]) + latencies.append(time.perf_counter() - text_started) + run_times.append(time.perf_counter() - run_started) + total_chars = sum(len(item["text"]) for item in items) + avg_run = statistics.mean(run_times) + return { + "texts_per_sec": len(items) / avg_run, + "chars_per_sec": total_chars / avg_run, + "latency_p50_ms": statistics.median(latencies) * 1000, + "latency_p95_ms": statistics.quantiles(latencies, n=20)[18] * 1000, + } + + +def spans(results, keep_ner: bool) -> set: + return { + (r.entity_type, r.start, r.end) + for r in results + if (r.entity_type in NER_ENTITIES) == keep_ner + } + + +def iou(a: tuple, b: tuple) -> float: + inter = max(0, min(a[2], b[2]) - max(a[1], b[1])) + union = max(a[2], b[2]) - min(a[1], b[1]) + return inter / union if union else 0.0 + + +def diff_ner(items, results_a, results_b, max_examples: int) -> dict: + """Per-entity-type agreement between two engines (span IoU >= 0.5).""" + per_type = defaultdict(lambda: {"a_total": 0, "b_total": 0, "matched": 0}) + examples = [] + for item, res_a, res_b in zip(items, results_a, results_b): + a = sorted(spans(res_a, keep_ner=True)) + b = sorted(spans(res_b, keep_ner=True)) + unmatched_b = set(b) + for span_a in a: + per_type[span_a[0]]["a_total"] += 1 + match = next( + (s for s in unmatched_b if s[0] == span_a[0] and iou(span_a, s) >= 0.5), None + ) + if match: + per_type[span_a[0]]["matched"] += 1 + unmatched_b.discard(match) + for span_b in b: + per_type[span_b[0]]["b_total"] += 1 + only_a = [s for s in a if not any(s[0] == t[0] and iou(s, t) >= 0.5 for t in b)] + only_b = sorted(unmatched_b) + if (only_a or only_b) and len(examples) < max_examples: + examples.append( + { + "text": item["text"], + "language": item["language"], + "only_a": [f"{t}[{s}:{e}]={item['text'][s:e]!r}" for t, s, e in only_a], + "only_b": [f"{t}[{s}:{e}]={item['text'][s:e]!r}" for t, s, e in only_b], + } + ) + return {"per_type": dict(per_type), "examples": examples} + + +def diff_regex(items, results_a, results_b) -> list: + """Non-NER results must be identical: same recognizers on both engines.""" + mismatches = [] + for item, res_a, res_b in zip(items, results_a, results_b): + a = spans(res_a, keep_ner=False) + b = spans(res_b, keep_ner=False) + if a != b: + mismatches.append({"text": item["text"], "only_a": sorted(a - b), "only_b": sorted(b - a)}) + return mismatches + + +def main() -> int: + args = parse_args() + items = json.loads(args.payload.read_text()) + engine_names = [e.strip() for e in args.engines.split(",") if e.strip()] + + report = {"payload": str(args.payload), "texts": len(items), "engines": {}} + results_by_engine = {} + for name in engine_names: + analyzer, build_secs = build(name, args) + stats = bench(analyzer, items, runs=args.runs, warmup=args.warmup) + stats["build_secs"] = build_secs + report["engines"][name] = stats + results_by_engine[name] = analyze_all(analyzer, items) + + exit_code = 0 + if set(engine_names) >= {"spacy", "gliner"}: + report["ner_diff"] = diff_ner( + items, results_by_engine["spacy"], results_by_engine["gliner"], args.max_examples + ) + regex_mismatches = diff_regex( + items, results_by_engine["spacy"], results_by_engine["gliner"] + ) + report["regex_mismatches"] = regex_mismatches + if regex_mismatches: + exit_code = 1 + + if args.json: + print(json.dumps(report, indent=2, default=str)) + return exit_code + + for name, stats in report["engines"].items(): + print(f"\n== {name} ==") + print(f" build: {stats['build_secs']:.1f}s") + print(f" throughput: {stats['texts_per_sec']:.2f} texts/s ({stats['chars_per_sec']:.0f} chars/s)") + print(f" latency: p50 {stats['latency_p50_ms']:.1f}ms p95 {stats['latency_p95_ms']:.1f}ms") + if "ner_diff" in report: + print("\n== NER parity (spacy=a vs gliner=b, span IoU>=0.5) ==") + for entity, counts in sorted(report["ner_diff"]["per_type"].items()): + print( + f" {entity:<10} spacy={counts['a_total']:<4} gliner={counts['b_total']:<4} " + f"matched={counts['matched']}" + ) + for example in report["ner_diff"]["examples"]: + print(f"\n [{example['language']}] {example['text']}") + if example["only_a"]: + print(f" spacy only: {', '.join(example['only_a'])}") + if example["only_b"]: + print(f" gliner only: {', '.join(example['only_b'])}") + if report["regex_mismatches"]: + print("\n!! REGEX MISMATCHES (wiring bug — engines must agree on non-NER):") + for mismatch in report["regex_mismatches"]: + print(f" {mismatch}") + else: + print("\n regex/checksum entities: identical across engines ✓") + return exit_code + + +if __name__ == "__main__": + raise SystemExit(main()) diff --git a/apps/pii/scripts/bench_payload.json b/apps/pii/scripts/bench_payload.json new file mode 100644 index 00000000000..fd0f207164d --- /dev/null +++ b/apps/pii/scripts/bench_payload.json @@ -0,0 +1,97 @@ +[ + { "text": "My name is John Smith and I live in Paris with my wife Marie.", "language": "en" }, + { + "text": "Dr. Angela Rodriguez will see you on March 14, 2026 at 3:30 PM in the Boston clinic.", + "language": "en" + }, + { + "text": "Contact Sarah O'Connor at sarah.oconnor@example.com or call (212) 555-0123.", + "language": "en" + }, + { "text": "The package ships to 45 Queen Street, Toronto, next Tuesday.", "language": "en" }, + { + "text": "Ahmed is a practicing Muslim from Egypt who moved to Berlin in 2019.", + "language": "en" + }, + { "text": "She is a Catholic Norwegian citizen born on 12/05/1988.", "language": "en" }, + { + "text": "Payment with card 4111111111111111 was declined yesterday at 10am.", + "language": "en" + }, + { + "text": "The vehicle VIN 1HGCM82633A004352 was registered to James Wilson in Ohio.", + "language": "en" + }, + { + "text": "His National Insurance number is AB123456C and he lives in Manchester.", + "language": "en" + }, + { + "text": "Meeting rescheduled: Friday, 9 January 2026, 14:00, with Priya Natarajan from Mumbai.", + "language": "en" + }, + { "text": "Send the invoice to accounts@acme.io before the end of Q1 2026.", "language": "en" }, + { + "text": "Klaus, a German engineer, and his Buddhist colleague Mei flew from Munich to Osaka.", + "language": "en" + }, + { + "text": "The Democrats and Republicans debated in Washington on election night.", + "language": "en" + }, + { + "text": "No PII here: the quarterly revenue grew 14% and margins held steady.", + "language": "en" + }, + { + "text": "Server request id a7f3k2m9x1q8w5z2b is not a VIN and not a person.", + "language": "en" + }, + { + "text": "Me llamo María García y vivo en Madrid desde el 3 de mayo de 2020.", + "language": "es" + }, + { "text": "Mi NIF es 12345678Z y mi correo es maria.garcia@ejemplo.es.", "language": "es" }, + { + "text": "El señor Javier Morales, ciudadano mexicano, llegó a Barcelona el lunes.", + "language": "es" + }, + { "text": "La reunión con Carmen será el 15 de junio de 2026 en Sevilla.", "language": "es" }, + { + "text": "Los musulmanes y los católicos convivieron durante siglos en Córdoba.", + "language": "es" + }, + { "text": "Mi chiamo Marco Rossi e abito a Roma vicino al Colosseo.", "language": "it" }, + { + "text": "Il codice fiscale di Maria Rossi è RSSMRA85T10A562S, nata il 10 dicembre 1985.", + "language": "it" + }, + { + "text": "Giulia Bianchi, cittadina italiana, si trasferì a Milano nel gennaio 2021.", + "language": "it" + }, + { + "text": "L'appuntamento con il dottor Ferrari è fissato per il 20 marzo 2026 a Torino.", + "language": "it" + }, + { "text": "Nazywam się Jan Kowalski i mieszkam w Warszawie od 2015 roku.", "language": "pl" }, + { + "text": "Mój numer PESEL to 44051401359, urodziłem się 14 maja 1944 w Krakowie.", + "language": "pl" + }, + { + "text": "Anna Nowak, obywatelka polska, spotka się z nami we wtorek 12 maja 2026.", + "language": "pl" + }, + { "text": "Katolicy i protestanci wspólnie świętowali w Gdańsku.", "language": "pl" }, + { + "text": "Nimeni on Matti Virtanen ja asun Helsingissä Töölön kaupunginosassa.", + "language": "fi" + }, + { + "text": "Henkilötunnukseni on 131052-308T ja synnyin Tampereella lokakuussa 1952.", + "language": "fi" + }, + { "text": "Liisa Korhonen muutti Ouluun maanantaina 5. tammikuuta 2026.", "language": "fi" }, + { "text": "Suomalaiset ja ruotsalaiset kilpailevat jääkiekossa joka vuosi.", "language": "fi" } +] diff --git a/apps/pii/server.py b/apps/pii/server.py index e2f7ad706d3..1f38031a79f 100644 --- a/apps/pii/server.py +++ b/apps/pii/server.py @@ -3,148 +3,52 @@ Constructs one warm AnalyzerEngine (multi-language NLP + a native check-digit VIN recognizer) and one AnonymizerEngine at startup, exposing stock-compatible endpoints so a single PRESIDIO_URL serves both. + +NER engine selection (see engines.py): +- PII_ENGINE=spacy (default): the 5 large spaCy models, unchanged behavior. +- PII_ENGINE=gliner: one multilingual GLiNER model for PERSON/LOCATION/NRP/ + DATE_TIME; requires the `gliner` image target (torch + gliner installed). + PII_DEVICE picks cpu/cuda (unset = auto-detect), PII_GLINER_MODEL overrides + the model id. The same code runs on CPU and GPU. Each uvicorn worker + (PII_WORKERS) loads its own GLiNER model copy — into GPU memory when on + cuda — so GPU deployments generally want PII_WORKERS=1 per GPU, unlike the + CPU/spacy path where workers scale with vCPUs. """ import logging +import os import time from typing import Any +from engines import build_gliner_analyzer, build_spacy_analyzer from fastapi import FastAPI -from presidio_analyzer import ( - AnalyzerEngine, - BatchAnalyzerEngine, - Pattern, - PatternRecognizer, - RecognizerResult, -) -from presidio_analyzer.nlp_engine import NlpEngineProvider -from presidio_analyzer.predefined_recognizers import ( - AuAbnRecognizer, - AuAcnRecognizer, - AuMedicareRecognizer, - AuTfnRecognizer, - EsNieRecognizer, - EsNifRecognizer, - FiPersonalIdentityCodeRecognizer, - InAadhaarRecognizer, - InPanRecognizer, - InPassportRecognizer, - InVehicleRegistrationRecognizer, - InVoterRecognizer, - ItDriverLicenseRecognizer, - ItFiscalCodeRecognizer, - ItIdentityCardRecognizer, - ItPassportRecognizer, - ItVatCodeRecognizer, - PlPeselRecognizer, - SgFinRecognizer, - SgUenRecognizer, - UkNinoRecognizer, -) +from presidio_analyzer import AnalyzerEngine, BatchAnalyzerEngine, RecognizerResult from presidio_anonymizer import AnonymizerEngine from presidio_anonymizer.entities import OperatorConfig from pydantic import BaseModel -# Languages served. Each needs its spaCy model installed in the image; the -# es/it/pl/fi predefined recognizers (ES_NIF, IT_FISCAL_CODE, PL_PESEL, ...) -# auto-load once their NLP engine is present. -NLP_CONFIGURATION = { - "nlp_engine_name": "spacy", - "models": [ - {"lang_code": "en", "model_name": "en_core_web_lg"}, - {"lang_code": "es", "model_name": "es_core_news_lg"}, - {"lang_code": "it", "model_name": "it_core_news_lg"}, - {"lang_code": "pl", "model_name": "pl_core_news_lg"}, - {"lang_code": "fi", "model_name": "fi_core_news_lg"}, - ], -} -SUPPORTED_LANGUAGES = [m["lang_code"] for m in NLP_CONFIGURATION["models"]] - -# Predefined recognizers Presidio ships but does NOT load into the default -# registry — they must be added explicitly. Each carries its own -# supported_language, so it fires under that language once its NLP model is -# loaded. en: UK/AU/IN/SG locale ids; es/it/pl/fi: national ids. -EXTRA_RECOGNIZERS = [ - UkNinoRecognizer, - AuAbnRecognizer, - AuAcnRecognizer, - AuTfnRecognizer, - AuMedicareRecognizer, - InPanRecognizer, - InAadhaarRecognizer, - InVehicleRegistrationRecognizer, - InVoterRecognizer, - InPassportRecognizer, - SgFinRecognizer, - SgUenRecognizer, - EsNifRecognizer, - EsNieRecognizer, - ItFiscalCodeRecognizer, - ItDriverLicenseRecognizer, - ItVatCodeRecognizer, - ItPassportRecognizer, - ItIdentityCardRecognizer, - PlPeselRecognizer, - FiPersonalIdentityCodeRecognizer, -] - - -class VinRecognizer(PatternRecognizer): - """VIN (17 chars, A-Z/0-9 excluding I/O/Q) with ISO 3779 check-digit - validation (position 9). Validation makes accidental matches on arbitrary - 17-char codes (request ids, SKUs, tokens) extremely unlikely. Some - non-North-American VINs omit the check digit and are skipped — an - intentional bias toward precision. - """ - - _TRANSLIT = { - **{str(d): d for d in range(10)}, - "A": 1, "B": 2, "C": 3, "D": 4, "E": 5, "F": 6, "G": 7, "H": 8, - "J": 1, "K": 2, "L": 3, "M": 4, "N": 5, "P": 7, "R": 9, - "S": 2, "T": 3, "U": 4, "V": 5, "W": 6, "X": 7, "Y": 8, "Z": 9, - } - _WEIGHTS = [8, 7, 6, 5, 4, 3, 2, 10, 0, 9, 8, 7, 6, 5, 4, 3, 2] +PII_ENGINE = os.environ.get("PII_ENGINE", "spacy") +if PII_ENGINE not in ("spacy", "gliner"): + raise ValueError(f"Invalid PII_ENGINE={PII_ENGINE!r}; expected 'spacy' or 'gliner'") +# Empty/unset -> None -> auto-detect (cuda when torch sees a GPU, else cpu). +PII_DEVICE = os.environ.get("PII_DEVICE") or None +PII_GLINER_MODEL = os.environ.get("PII_GLINER_MODEL", "urchade/gliner_multi_pii-v1") - def validate_result(self, pattern_text: str): - vin = pattern_text.upper() - if len(vin) != 17: - return False - try: - total = sum(self._TRANSLIT[c] * w for c, w in zip(vin, self._WEIGHTS)) - except KeyError: - return False - check = total % 11 - expected = "X" if check == 10 else str(check) - return vin[8] == expected +# Propagates to uvicorn's root handler, so timing lands in the container log stream. +logger = logging.getLogger("sim.pii") def build_analyzer() -> AnalyzerEngine: - nlp_engine = NlpEngineProvider(nlp_configuration=NLP_CONFIGURATION).create_engine() - analyzer = AnalyzerEngine(nlp_engine=nlp_engine, supported_languages=SUPPORTED_LANGUAGES) - # VIN is language-agnostic, so register it under every served language — - # a recognizer only fires for the language the caller routes to. - vin_pattern = Pattern(name="vin", regex=r"\b[A-HJ-NPR-Z0-9]{17}\b", score=0.7) - for language in SUPPORTED_LANGUAGES: - analyzer.registry.add_recognizer( - VinRecognizer( - supported_entity="VIN", - patterns=[vin_pattern], - context=["vin", "vehicle", "chassis"], - supported_language=language, - ) - ) - for recognizer_cls in EXTRA_RECOGNIZERS: - analyzer.registry.add_recognizer(recognizer_cls()) - return analyzer + if PII_ENGINE == "gliner": + return build_gliner_analyzer(model_name=PII_GLINER_MODEL, device=PII_DEVICE) + return build_spacy_analyzer() +logger.info("building analyzer engine=%s device=%s", PII_ENGINE, PII_DEVICE or "auto") analyzer = build_analyzer() batch_analyzer = BatchAnalyzerEngine(analyzer_engine=analyzer) anonymizer = AnonymizerEngine() -# Propagates to uvicorn's root handler, so timing lands in the container log stream. -logger = logging.getLogger("sim.pii") - app = FastAPI(title="Sim Presidio", docs_url=None, redoc_url=None) diff --git a/apps/pii/tests/conftest.py b/apps/pii/tests/conftest.py new file mode 100644 index 00000000000..c9787a309d9 --- /dev/null +++ b/apps/pii/tests/conftest.py @@ -0,0 +1,5 @@ +import sys +from pathlib import Path + +# server.py / engines.py live one level up (repo: apps/pii, image: /app). +sys.path.insert(0, str(Path(__file__).resolve().parent.parent)) diff --git a/apps/pii/tests/test_engines.py b/apps/pii/tests/test_engines.py new file mode 100644 index 00000000000..82a485a8265 --- /dev/null +++ b/apps/pii/tests/test_engines.py @@ -0,0 +1,105 @@ +"""Unit tests for engines.py — no models, no downloads, no network. + +Run: pip install -r requirements.txt -r requirements-dev.txt && python -m pytest tests +""" + +import importlib.util +import os +import subprocess +import sys +from pathlib import Path + +import pytest +from presidio_analyzer.predefined_recognizers.ner import gliner_recognizer + +import engines + +PII_DIR = Path(__file__).resolve().parent.parent + + +class FakeModel: + def __init__(self): + self.seen_labels: list[list[str]] = [] + + def predict_entities(self, text, labels, flat_ner=True, threshold=0.3, multi_label=False): + self.seen_labels.append(list(labels)) + return [{"label": "person", "score": 0.92, "start": 0, "end": 4, "text": text[0:4]}] + + +class FakeGLiNER: + calls = 0 + + @classmethod + def from_pretrained(cls, model_name, **kwargs): + cls.calls += 1 + return FakeModel() + + +@pytest.fixture +def fake_gliner(monkeypatch): + monkeypatch.setattr(gliner_recognizer, "GLiNER", FakeGLiNER) + engines.SharedModelGLiNERRecognizer._shared_models.clear() + FakeGLiNER.calls = 0 + yield FakeGLiNER + engines.SharedModelGLiNERRecognizer._shared_models.clear() + + +def make_recognizer(language: str): + return engines.SharedModelGLiNERRecognizer( + entity_mapping=engines.GLINER_ENTITY_MAPPING, + model_name="fake/model", + map_location="cpu", + supported_language=language, + ) + + +def test_invalid_pii_engine_fails_import(): + result = subprocess.run( + [sys.executable, "-c", "import server"], + cwd=PII_DIR, + env={**os.environ, "PII_ENGINE": "bogus"}, + capture_output=True, + text=True, + ) + assert result.returncode != 0 + assert "Invalid PII_ENGINE" in result.stderr + + +@pytest.mark.skipif( + importlib.util.find_spec("gliner") is not None, + reason="fail-fast path only exists when gliner is not installed", +) +def test_build_gliner_analyzer_fails_fast_without_gliner(): + with pytest.raises(RuntimeError, match="gliner image variant"): + engines.build_gliner_analyzer(model_name="fake/model", device="cpu") + + +def test_shared_model_loads_once_across_languages(fake_gliner): + first = make_recognizer("en") + second = make_recognizer("es") + assert fake_gliner.calls == 1 + assert first.gliner is second.gliner + + +def test_analyze_never_prompts_gliner_with_foreign_entities(fake_gliner): + recognizer = make_recognizer("en") + all_supported = ["PERSON", "LOCATION", "NRP", "DATE_TIME", "CREDIT_CARD", "VIN", "ES_NIF"] + results = recognizer.analyze("John went home", entities=all_supported) + for labels in recognizer.gliner.seen_labels: + assert set(labels) <= set(engines.GLINER_ENTITY_MAPPING) + assert results and results[0].entity_type == "PERSON" + + +def test_analyze_skips_inference_when_no_owned_entity_requested(fake_gliner): + recognizer = make_recognizer("en") + assert recognizer.analyze("4111111111111111", entities=["CREDIT_CARD"]) == [] + assert recognizer.gliner.seen_labels == [] + + +def test_entity_mapping_targets_exactly_the_ner_entities(): + assert set(engines.GLINER_ENTITY_MAPPING.values()) == { + "PERSON", + "LOCATION", + "NRP", + "DATE_TIME", + } diff --git a/apps/pii/tests/test_integration.py b/apps/pii/tests/test_integration.py new file mode 100644 index 00000000000..45d15ddd258 --- /dev/null +++ b/apps/pii/tests/test_integration.py @@ -0,0 +1,102 @@ +"""Integration tests — exercise the real engines end-to-end via the FastAPI app. + +Requires the models present, so run inside the built images (gated behind +RUN_PII_INTEGRATION to keep plain `pytest` runs model-free): + + # spacy regression (default engine) + docker run --rm -e RUN_PII_INTEGRATION=1 sim-pii:gliner python -m pytest tests + + # gliner engine + docker run --rm -e RUN_PII_INTEGRATION=1 -e PII_ENGINE=gliner sim-pii:gliner \ + python -m pytest tests/test_integration.py + +The suite adapts to PII_ENGINE: shared assertions always run, engine-specific +ones only for the active engine. +""" + +import os + +import pytest + +if not os.environ.get("RUN_PII_INTEGRATION"): + pytest.skip( + "integration tests need the built image (RUN_PII_INTEGRATION=1)", + allow_module_level=True, + ) + +from fastapi.testclient import TestClient + +import server + +ENGINE = server.PII_ENGINE +client = TestClient(server.app) + + +def redact_batch(texts, language="en"): + response = client.post("/redact_batch", json={"texts": texts, "language": language}) + assert response.status_code == 200 + return response.json()["texts"] + + +def test_health(): + assert client.get("/health").json() == {"status": "ok"} + + +def test_masks_person_and_email(): + [masked] = redact_batch(["My name is John Smith, email john.smith@example.com."]) + assert "" in masked + assert "" in masked + assert "John Smith" not in masked + assert "john.smith@example.com" not in masked + + +def test_masks_location_and_phone(): + [masked] = redact_batch(["I live in Paris, call me at (212) 555-0123."]) + assert "" in masked + assert "" in masked + assert "Paris" not in masked + + +def test_regex_recognizers_fire_in_non_english_languages(): + [masked] = redact_batch(["Mi NIF es 12345678Z."], language="es") + assert "" in masked + # On the spacy engine the it_core_news_lg NER tags the fiscal code as + # ORGANIZATION and outscores the pattern recognizer, so only assert the + # value is masked; the exact label is checked on the gliner engine where + # spaCy NER can't compete. + [masked] = redact_batch(["Il codice fiscale è RSSMRA85T10A562S."], language="it") + assert "RSSMRA85T10A562S" not in masked + if ENGINE == "gliner": + assert "" in masked + + +def test_vin_checksum_recognizer_fires(): + [masked] = redact_batch(["The car VIN is 1HGCM82633A004352."]) + assert "" in masked + + +def test_no_pii_passes_through_unchanged(): + # NB: "Quarterly" would be tagged DATE_TIME by the spacy engine — keep + # this text free of anything either engine considers an entity. + text = "Revenue grew and margins held steady." + assert redact_batch([text]) == [text] + + +@pytest.mark.skipif(ENGINE != "gliner", reason="gliner-only wiring assertions") +def test_gliner_registry_has_no_spacy_recognizer(): + names = {r.name for r in server.analyzer.registry.recognizers} + assert "SpacyRecognizer" not in names + assert "GLiNERRecognizer" in names + + +@pytest.mark.skipif(ENGINE != "gliner", reason="gliner-only wiring assertions") +def test_gliner_supported_entities_keep_ner_types(): + supported = set(server.analyzer.get_supported_entities("en")) + assert {"PERSON", "LOCATION", "NRP", "DATE_TIME"} <= supported + + +@pytest.mark.skipif(ENGINE != "spacy", reason="spacy-only wiring assertions") +def test_spacy_registry_unchanged(): + names = {r.name for r in server.analyzer.registry.recognizers} + assert "SpacyRecognizer" in names + assert "GLiNERRecognizer" not in names diff --git a/docker/pii.Dockerfile b/docker/pii.Dockerfile index 73f56647909..b04cbadb234 100644 --- a/docker/pii.Dockerfile +++ b/docker/pii.Dockerfile @@ -1,5 +1,17 @@ # ======================================== # Combined Presidio service (analyzer + anonymizer) on a single port (5001) +# +# Targets (docker builds the LAST stage when no --target is given): +# (default) / spacy : lean image, spaCy NER only — what self-hosters get. +# gliner : superset image (torch CPU + gliner + baked GLiNER +# model + small spaCy models). Both PII_ENGINE=spacy +# and PII_ENGINE=gliner work in it. +# gliner-gpu : scaffold for the EC2-GPU fleet — same layout, CUDA +# torch wheels (bundle their own CUDA libs; host only +# needs the nvidia container runtime). Not built in CI. +# +# Source files are COPY'd last in each terminal stage so code edits never +# re-download deps or models. # ======================================== FROM python:3.12-slim-bookworm AS base @@ -17,9 +29,26 @@ COPY apps/pii/requirements.txt ./requirements.txt RUN --mount=type=cache,target=/root/.cache/pip \ pip install -r requirements.txt +RUN groupadd -g 1001 pii && \ + useradd -u 1001 -g pii pii && \ + chown -R pii:pii /app + +# Listen on 5001. Runs as its own ECS service (separate task), reached via PII_URL; +# 5001 avoids colliding with the app's 3000 in local/compose runs on one host. +EXPOSE 5001 + +# start-period is generous: five large spaCy models load at import before +# /health responds. Tune against measured cold-start once built. +HEALTHCHECK --interval=30s --timeout=5s --start-period=180s --retries=3 \ + CMD curl -fsS http://localhost:5001/health || exit 1 + +CMD ["uvicorn", "server:app", "--host", "0.0.0.0", "--port", "5001"] + # Pinned spaCy models (en + es/it/pl/fi, ~2.2GB total). Downloaded with # retries/resume — the large wheels truncate on flaky networks if pip fetches -# the URLs directly. +# the URLs directly. Shared by every terminal image so PII_ENGINE=spacy works +# everywhere (the gliner opt-in stays reversible without an image swap). +FROM base AS spacy-models ARG SPACY_MODELS="en_core_web_lg-3.8.0 es_core_news_lg-3.8.0 it_core_news_lg-3.8.0 pl_core_news_lg-3.8.0 fi_core_news_lg-3.8.0" RUN --mount=type=cache,target=/root/.cache/pip \ for model in ${SPACY_MODELS}; do \ @@ -31,20 +60,90 @@ RUN --mount=type=cache,target=/root/.cache/pip \ pip install /tmp/*.whl && \ rm /tmp/*.whl -COPY apps/pii/server.py ./server.py +# --- GLiNER (CPU) ------------------------------------------------------------ +FROM spacy-models AS gliner + +# torch pinned here (not requirements-gliner.txt) because the CPU and CUDA +# targets install the same version from different wheel indexes. 2.11.0 is the +# newest release published on both the cpu and cu128 indexes for py312. +ARG TORCH_VERSION=2.11.0 +RUN --mount=type=cache,target=/root/.cache/pip \ + pip install torch==${TORCH_VERSION} --index-url https://download.pytorch.org/whl/cpu + +COPY apps/pii/requirements-gliner.txt ./requirements-gliner.txt +RUN --mount=type=cache,target=/root/.cache/pip \ + pip install -r requirements-gliner.txt + +# Small spaCy models (~60MB total) give the gliner engine tokenization + +# lemmas for the regex recognizers; GLiNER does the NER (see engines.py). +ARG SPACY_SM_MODELS="en_core_web_sm-3.8.0 es_core_news_sm-3.8.0 it_core_news_sm-3.8.0 pl_core_news_sm-3.8.0 fi_core_news_sm-3.8.0" +RUN --mount=type=cache,target=/root/.cache/pip \ + for model in ${SPACY_SM_MODELS}; do \ + whl="${model}-py3-none-any.whl"; \ + curl -fL --retry 5 --retry-delay 5 --retry-all-errors -C - \ + -o "/tmp/${whl}" \ + "https://github.com/explosion/spacy-models/releases/download/${model}/${whl}" || exit 1; \ + done && \ + pip install /tmp/*.whl && \ + rm /tmp/*.whl + +# Bake the GLiNER weights at build time (cached layer) so startup never +# touches the network. HF_HUB_OFFLINE makes a missing/overridden +# PII_GLINER_MODEL fail fast at startup instead of silently downloading. +ENV HF_HOME=/opt/hf-cache +ARG GLINER_MODEL=urchade/gliner_multi_pii-v1 +RUN python -c "from gliner import GLiNER; GLiNER.from_pretrained('${GLINER_MODEL}')" && \ + chmod -R a+rX /opt/hf-cache +ENV HF_HUB_OFFLINE=1 + +# Bench + tests ride along only in this image — it's the only one with both +# engines installed (see apps/pii/scripts/bench_engines.py). pytest/httpx are +# baked in so the documented `docker run ... python -m pytest tests` works +# as-is (the runtime user has no writable HOME for pip install --user). +COPY apps/pii/requirements-dev.txt ./requirements-dev.txt +RUN --mount=type=cache,target=/root/.cache/pip \ + pip install -r requirements-dev.txt + +COPY --chown=pii:pii apps/pii/server.py apps/pii/engines.py ./ +COPY --chown=pii:pii apps/pii/scripts ./scripts +COPY --chown=pii:pii apps/pii/tests ./tests -RUN groupadd -g 1001 pii && \ - useradd -u 1001 -g pii pii && \ - chown -R pii:pii /app USER pii -# Listen on 5001. Runs as its own ECS service (separate task), reached via PII_URL; -# 5001 avoids colliding with the app's 3000 in local/compose runs on one host. -EXPOSE 5001 +# --- GLiNER (CUDA) scaffold — wired for the GPU fleet follow-up -------------- +FROM spacy-models AS gliner-gpu -# start-period is generous: five large spaCy models load at import before -# /health responds. Tune against measured cold-start once built. -HEALTHCHECK --interval=30s --timeout=5s --start-period=180s --retries=3 \ - CMD curl -fsS http://localhost:5001/health || exit 1 +ARG TORCH_VERSION=2.11.0 +RUN --mount=type=cache,target=/root/.cache/pip \ + pip install torch==${TORCH_VERSION} --index-url https://download.pytorch.org/whl/cu128 -CMD ["uvicorn", "server:app", "--host", "0.0.0.0", "--port", "5001"] +COPY apps/pii/requirements-gliner.txt ./requirements-gliner.txt +RUN --mount=type=cache,target=/root/.cache/pip \ + pip install -r requirements-gliner.txt + +ARG SPACY_SM_MODELS="en_core_web_sm-3.8.0 es_core_news_sm-3.8.0 it_core_news_sm-3.8.0 pl_core_news_sm-3.8.0 fi_core_news_sm-3.8.0" +RUN --mount=type=cache,target=/root/.cache/pip \ + for model in ${SPACY_SM_MODELS}; do \ + whl="${model}-py3-none-any.whl"; \ + curl -fL --retry 5 --retry-delay 5 --retry-all-errors -C - \ + -o "/tmp/${whl}" \ + "https://github.com/explosion/spacy-models/releases/download/${model}/${whl}" || exit 1; \ + done && \ + pip install /tmp/*.whl && \ + rm /tmp/*.whl + +ENV HF_HOME=/opt/hf-cache +ARG GLINER_MODEL=urchade/gliner_multi_pii-v1 +RUN python -c "from gliner import GLiNER; GLiNER.from_pretrained('${GLINER_MODEL}')" && \ + chmod -R a+rX /opt/hf-cache +ENV HF_HUB_OFFLINE=1 + +COPY --chown=pii:pii apps/pii/server.py apps/pii/engines.py ./ +COPY --chown=pii:pii apps/pii/scripts ./scripts + +USER pii + +# --- DEFAULT (last stage): the lean spaCy image, content-equivalent to today - +FROM spacy-models AS spacy +COPY --chown=pii:pii apps/pii/server.py apps/pii/engines.py ./ +USER pii diff --git a/helm/sim/templates/deployment-pii.yaml b/helm/sim/templates/deployment-pii.yaml index e6c7dc206f0..2dd4b8b22c6 100644 --- a/helm/sim/templates/deployment-pii.yaml +++ b/helm/sim/templates/deployment-pii.yaml @@ -44,16 +44,20 @@ spec: - name: http containerPort: {{ .Values.pii.service.targetPort }} protocol: TCP - {{- if or .Values.pii.env .Values.extraEnvVars }} env: - {{- range $key, $value := .Values.pii.env }} + - name: PII_ENGINE + value: {{ .Values.pii.engine | quote }} + {{- if .Values.pii.device }} + - name: PII_DEVICE + value: {{ .Values.pii.device | quote }} + {{- end }} + {{- range $key, $value := omit (.Values.pii.env | default dict) "PII_ENGINE" "PII_DEVICE" }} - name: {{ $key }} value: {{ $value | quote }} {{- end }} {{- with .Values.extraEnvVars }} {{- toYaml . | nindent 12 }} {{- end }} - {{- end }} {{- if .Values.pii.startupProbe }} startupProbe: {{- toYaml .Values.pii.startupProbe | nindent 12 }} diff --git a/helm/sim/tests/pii_test.yaml b/helm/sim/tests/pii_test.yaml index 9d89b625ac5..3a416237207 100644 --- a/helm/sim/tests/pii_test.yaml +++ b/helm/sim/tests/pii_test.yaml @@ -31,6 +31,65 @@ tests: path: spec.template.spec.containers[0].ports[0].containerPort value: 5001 + - it: pii pod defaults to the spacy engine + template: deployment-pii.yaml + set: + pii.enabled: true + asserts: + - contains: + path: spec.template.spec.containers[0].env + content: + name: PII_ENGINE + value: "spacy" + - notContains: + path: spec.template.spec.containers[0].env + content: + name: PII_DEVICE + value: "cpu" + + - it: pii.engine and pii.device render through to the container env + template: deployment-pii.yaml + set: + pii.enabled: true + pii.engine: gliner + pii.device: cuda + asserts: + - contains: + path: spec.template.spec.containers[0].env + content: + name: PII_ENGINE + value: "gliner" + - contains: + path: spec.template.spec.containers[0].env + content: + name: PII_DEVICE + value: "cuda" + + - it: user-set pii.env cannot override the chart-owned engine keys + template: deployment-pii.yaml + set: + pii.enabled: true + pii.env: + PII_ENGINE: evil + PII_DEVICE: evil + PII_GLINER_MODEL: acme/custom-model + asserts: + - notContains: + path: spec.template.spec.containers[0].env + content: + name: PII_ENGINE + value: "evil" + - notContains: + path: spec.template.spec.containers[0].env + content: + name: PII_DEVICE + value: "evil" + - contains: + path: spec.template.spec.containers[0].env + content: + name: PII_GLINER_MODEL + value: "acme/custom-model" + - it: app pod gets chart-computed PII_URL pointing at the in-cluster service template: deployment-app.yaml set: diff --git a/helm/sim/values.schema.json b/helm/sim/values.schema.json index e20ec2026f2..71ec431f6d5 100644 --- a/helm/sim/values.schema.json +++ b/helm/sim/values.schema.json @@ -711,6 +711,15 @@ "type": "boolean", "description": "Enable the Presidio PII redaction service" }, + "engine": { + "type": "string", + "enum": ["spacy", "gliner"], + "description": "NER engine; gliner requires the -gliner image variant" + }, + "device": { + "type": "string", + "description": "Torch device for the gliner engine (cpu, cuda, cuda:N); empty auto-detects" + }, "replicaCount": { "type": "integer", "minimum": 1, diff --git a/helm/sim/values.yaml b/helm/sim/values.yaml index 91ae16747fe..52d320f89d5 100644 --- a/helm/sim/values.yaml +++ b/helm/sim/values.yaml @@ -824,6 +824,19 @@ pii: digest: "" # sha256: pin overrides tag pullPolicy: IfNotPresent + # NER engine: "spacy" (default) or "gliner" (opt-in zero-shot transformer NER + # for PERSON/LOCATION/NRP/DATE_TIME; regex/checksum recognizers are identical + # on both engines). gliner requires the gliner image variant — the `-gliner` + # published tag (e.g. image.tag: latest-gliner) or a local + # `docker build --target gliner -f docker/pii.Dockerfile` — and fails fast at + # startup on the default image. + engine: "spacy" + + # Torch device for the gliner engine: "cpu", "cuda", or "cuda:N". Empty + # auto-detects (cuda when a GPU is visible, else cpu). GPU resource requests + # (nvidia.com/gpu) are not wired yet — GPU scheduling is an infra follow-up. + device: "" + # Number of replicas replicaCount: 1 From bff762688a6874e04762dd687127676b52bc2db9 Mon Sep 17 00:00:00 2001 From: Theodore Li Date: Tue, 7 Jul 2026 17:20:06 -0700 Subject: [PATCH 2/2] =?UTF-8?q?refactor(pii):=20ship=20both=20engines=20in?= =?UTF-8?q?=20one=20image=20=E2=80=94=20engine=20is=20a=20pure=20env=20fli?= =?UTF-8?q?p?= MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit Collapse the gliner build target into the single pii image: spaCy lg models, torch (CPU), gliner, and the baked GLiNER weights all ship in it, so PII_ENGINE switches engines with no image swap and no tag matrix. CI reverts to the single pii build (no -gliner tags). The GPU variant becomes the same Dockerfile built with --build-arg TORCH_INDEX_URL=.../cu128. Image grows ~6.1GB -> ~9.6GB. Co-Authored-By: Claude Fable 5 Claude-Session: https://claude.ai/code/session_01Up3F97mjCH9HCj1pX4J8VJ --- .github/workflows/ci.yml | 38 ++------- apps/pii/engines.py | 15 ++-- apps/pii/scripts/bench_engines.py | 6 +- apps/pii/server.py | 2 +- apps/pii/tests/test_engines.py | 2 +- apps/pii/tests/test_integration.py | 4 +- docker/pii.Dockerfile | 121 ++++++++++------------------- helm/sim/values.schema.json | 2 +- helm/sim/values.yaml | 7 +- 9 files changed, 65 insertions(+), 132 deletions(-) diff --git a/.github/workflows/ci.yml b/.github/workflows/ci.yml index d2137803c05..25d8348d8c4 100644 --- a/.github/workflows/ci.yml +++ b/.github/workflows/ci.yml @@ -90,12 +90,6 @@ jobs: ecr_repo_secret: ECR_REALTIME - dockerfile: ./docker/pii.Dockerfile ecr_repo_secret: ECR_PII - # Opt-in GLiNER variant of the pii image (torch + gliner + baked - # model). Deployed by setting PII_ENGINE=gliner on the pii service. - - dockerfile: ./docker/pii.Dockerfile - ecr_repo_secret: ECR_PII - target: gliner - tag_suffix: -gliner steps: - name: Checkout code uses: actions/checkout@df4cb1c069e1874edd31b4311f1884172cec0e10 # v6 @@ -132,8 +126,7 @@ jobs: file: ${{ matrix.dockerfile }} platforms: linux/amd64 push: true - tags: ${{ steps.login-ecr.outputs.registry }}/${{ steps.ecr-repo.outputs.name }}:dev${{ matrix.tag_suffix || '' }} - target: ${{ matrix.target || '' }} + tags: ${{ steps.login-ecr.outputs.registry }}/${{ steps.ecr-repo.outputs.name }}:dev provenance: false sbom: false @@ -208,13 +201,6 @@ jobs: - dockerfile: ./docker/pii.Dockerfile ghcr_image: ghcr.io/simstudioai/pii ecr_repo_secret: ECR_PII - # Opt-in GLiNER variant of the pii image (torch + gliner + baked - # model). amd64 only — the deploy fleet is amd64; no arm64/manifest. - - dockerfile: ./docker/pii.Dockerfile - ghcr_image: ghcr.io/simstudioai/pii - ecr_repo_secret: ECR_PII - target: gliner - tag_suffix: -gliner steps: - name: Checkout code uses: actions/checkout@df4cb1c069e1874edd31b4311f1884172cec0e10 # v6 @@ -258,35 +244,26 @@ jobs: ECR_REGISTRY="${{ steps.login-ecr.outputs.registry }}" ECR_REPO="${{ steps.ecr-repo.outputs.name }}" GHCR_IMAGE="${{ matrix.ghcr_image }}" - TAG_SUFFIX="${{ matrix.tag_suffix }}" - - # Suffixed variants (e.g. pii -gliner) are amd64-only and skip the - # manifest merge, so their GHCR tags are plain — no -amd64 alias. - if [ -n "$TAG_SUFFIX" ]; then - ARCH_SUFFIX="" - else - ARCH_SUFFIX="-amd64" - fi if [ "${{ github.ref }}" = "refs/heads/main" ]; then - ECR_TAG="latest${TAG_SUFFIX}" + ECR_TAG="latest" else - ECR_TAG="staging${TAG_SUFFIX}" + ECR_TAG="staging" fi ECR_IMAGE="${ECR_REGISTRY}/${ECR_REPO}:${ECR_TAG}" TAGS="${ECR_IMAGE}" if [ "${{ github.ref }}" = "refs/heads/main" ] && [ -n "$GHCR_IMAGE" ]; then - GHCR_AMD64="${GHCR_IMAGE}:latest${TAG_SUFFIX}${ARCH_SUFFIX}" - GHCR_SHA="${GHCR_IMAGE}:${{ github.sha }}${TAG_SUFFIX}${ARCH_SUFFIX}" + GHCR_AMD64="${GHCR_IMAGE}:latest-amd64" + GHCR_SHA="${GHCR_IMAGE}:${{ github.sha }}-amd64" TAGS="${TAGS},$GHCR_AMD64,$GHCR_SHA" if [ "${{ needs.detect-version.outputs.is_release }}" = "true" ]; then VERSION="${{ needs.detect-version.outputs.version }}" - GHCR_VERSION="${GHCR_IMAGE}:${VERSION}${TAG_SUFFIX}${ARCH_SUFFIX}" + GHCR_VERSION="${GHCR_IMAGE}:${VERSION}-amd64" TAGS="${TAGS},$GHCR_VERSION" - echo "📦 Adding version tag: ${VERSION}${TAG_SUFFIX}${ARCH_SUFFIX}" + echo "📦 Adding version tag: ${VERSION}-amd64" fi fi @@ -300,7 +277,6 @@ jobs: platforms: linux/amd64 push: true tags: ${{ steps.meta.outputs.tags }} - target: ${{ matrix.target || '' }} provenance: false sbom: false diff --git a/apps/pii/engines.py b/apps/pii/engines.py index 8af7184d26c..c8edf982ff1 100644 --- a/apps/pii/engines.py +++ b/apps/pii/engines.py @@ -225,14 +225,15 @@ def build_gliner_analyzer(model_name: str, device: str | None) -> AnalyzerEngine :param device: torch device ("cpu", "cuda", "cuda:0"); None auto-detects via Presidio's device_detector (cuda when available, else cpu). """ - # Fail fast with an actionable message on the lean image. Without these - # checks Presidio would try to pip-download the missing spaCy models at - # startup (a silent network fallback that dies with an unrelated pip - # permission error), and the gliner ImportError would surface only later. + # Fail fast with an actionable message when gliner deps are missing (e.g. + # a custom-built image without them). Without these checks Presidio would + # try to pip-download the missing spaCy models at startup (a silent + # network fallback that dies with an unrelated pip permission error), and + # the gliner ImportError would surface only later. if importlib.util.find_spec("gliner") is None: raise RuntimeError( - "PII_ENGINE=gliner requires the gliner image variant " - "(docker build --target gliner); the gliner package is not installed" + "PII_ENGINE=gliner but the gliner package is not installed; " + "use the stock pii image (docker/pii.Dockerfile ships torch + gliner)" ) missing = [ m["model_name"] @@ -242,7 +243,7 @@ def build_gliner_analyzer(model_name: str, device: str | None) -> AnalyzerEngine if missing: raise RuntimeError( f"PII_ENGINE=gliner needs spaCy models {missing}; " - "use the gliner image variant (docker build --target gliner)" + "use the stock pii image (docker/pii.Dockerfile ships them)" ) nlp_engine = NlpEngineProvider(nlp_configuration=GLINER_NLP_CONFIGURATION).create_engine() analyzer = AnalyzerEngine(nlp_engine=nlp_engine, supported_languages=SUPPORTED_LANGUAGES) diff --git a/apps/pii/scripts/bench_engines.py b/apps/pii/scripts/bench_engines.py index 8d902fe4d22..b7ea70ca764 100644 --- a/apps/pii/scripts/bench_engines.py +++ b/apps/pii/scripts/bench_engines.py @@ -7,10 +7,10 @@ register the same recognizers — so any mismatch there is a wiring bug and the script exits non-zero. -Meant to run inside the gliner image (the only one with both engines): +Meant to run inside the pii image (both engines ship in it): - docker run --rm sim-pii:gliner python scripts/bench_engines.py - docker run --rm -v $PWD/texts.json:/data.json sim-pii:gliner \\ + docker run --rm python scripts/bench_engines.py + docker run --rm -v $PWD/texts.json:/data.json \\ python scripts/bench_engines.py --payload /data.json Payload format: JSON list of {"text": str, "language": str} objects. diff --git a/apps/pii/server.py b/apps/pii/server.py index 1f38031a79f..3f59cc540aa 100644 --- a/apps/pii/server.py +++ b/apps/pii/server.py @@ -7,7 +7,7 @@ NER engine selection (see engines.py): - PII_ENGINE=spacy (default): the 5 large spaCy models, unchanged behavior. - PII_ENGINE=gliner: one multilingual GLiNER model for PERSON/LOCATION/NRP/ - DATE_TIME; requires the `gliner` image target (torch + gliner installed). + DATE_TIME. The stock image ships both engines, so this is a pure env flip. PII_DEVICE picks cpu/cuda (unset = auto-detect), PII_GLINER_MODEL overrides the model id. The same code runs on CPU and GPU. Each uvicorn worker (PII_WORKERS) loads its own GLiNER model copy — into GPU memory when on diff --git a/apps/pii/tests/test_engines.py b/apps/pii/tests/test_engines.py index 82a485a8265..1e6c1b8840c 100644 --- a/apps/pii/tests/test_engines.py +++ b/apps/pii/tests/test_engines.py @@ -70,7 +70,7 @@ def test_invalid_pii_engine_fails_import(): reason="fail-fast path only exists when gliner is not installed", ) def test_build_gliner_analyzer_fails_fast_without_gliner(): - with pytest.raises(RuntimeError, match="gliner image variant"): + with pytest.raises(RuntimeError, match="gliner package is not installed"): engines.build_gliner_analyzer(model_name="fake/model", device="cpu") diff --git a/apps/pii/tests/test_integration.py b/apps/pii/tests/test_integration.py index 45d15ddd258..40c97975754 100644 --- a/apps/pii/tests/test_integration.py +++ b/apps/pii/tests/test_integration.py @@ -4,10 +4,10 @@ RUN_PII_INTEGRATION to keep plain `pytest` runs model-free): # spacy regression (default engine) - docker run --rm -e RUN_PII_INTEGRATION=1 sim-pii:gliner python -m pytest tests + docker run --rm -e RUN_PII_INTEGRATION=1 python -m pytest tests # gliner engine - docker run --rm -e RUN_PII_INTEGRATION=1 -e PII_ENGINE=gliner sim-pii:gliner \ + docker run --rm -e RUN_PII_INTEGRATION=1 -e PII_ENGINE=gliner \ python -m pytest tests/test_integration.py The suite adapts to PII_ENGINE: shared assertions always run, engine-specific diff --git a/docker/pii.Dockerfile b/docker/pii.Dockerfile index fa532d01a71..5d0fedb542c 100644 --- a/docker/pii.Dockerfile +++ b/docker/pii.Dockerfile @@ -1,17 +1,17 @@ # ======================================== # Combined Presidio service (analyzer + anonymizer) on a single port (5001) # -# Targets (docker builds the LAST stage when no --target is given): -# (default) / spacy : lean image, spaCy NER only — what self-hosters get. -# gliner : superset image (torch CPU + gliner + baked GLiNER -# model + small spaCy models). Both PII_ENGINE=spacy -# and PII_ENGINE=gliner work in it. -# gliner-gpu : scaffold for the EC2-GPU fleet — same layout, CUDA -# torch wheels (bundle their own CUDA libs; host only -# needs the nvidia container runtime). Not built in CI. +# ONE image serves both NER engines — the engine is a pure runtime choice via +# PII_ENGINE (spacy default | gliner). spaCy large models, torch (CPU), the +# gliner package, and the baked GLiNER weights all ship in it, so flipping +# engines never requires an image swap. # -# Source files are COPY'd last in each terminal stage so code edits never -# re-download deps or models. +# GPU variant (EC2-GPU fleet follow-up): same Dockerfile, CUDA torch wheels — +# docker build --build-arg TORCH_INDEX_URL=https://download.pytorch.org/whl/cu128 ... +# (torch CUDA wheels bundle their own CUDA libs; the host only needs the +# nvidia container runtime.) +# +# Source files are COPY'd last so code edits never re-download deps or models. # ======================================== FROM python:3.12-slim-bookworm AS base @@ -29,35 +29,9 @@ COPY apps/pii/requirements.txt ./requirements.txt RUN --mount=type=cache,target=/root/.cache/pip \ pip install -r requirements.txt -RUN groupadd -g 1001 pii && \ - useradd -u 1001 -g pii pii && \ - chown -R pii:pii /app - -# Listen on 5001. Runs as its own ECS service (separate task), reached via PII_URL; -# 5001 avoids colliding with the app's 3000 in local/compose runs on one host. -EXPOSE 5001 - -# start-period covers the model cold start. With PII_WORKERS>1 each worker loads -# the five spaCy models independently and in parallel, so allow generous headroom -# (memory-bandwidth contention stretches the wall-time beyond the single-worker case). -HEALTHCHECK --interval=30s --timeout=5s --start-period=300s --retries=3 \ - CMD curl -fsS http://localhost:5001/health || exit 1 - -# 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. -# NB for the gliner engine: EACH worker loads its own GLiNER model copy (into GPU -# memory when on cuda), so GPU deployments generally want PII_WORKERS=1 per GPU. -CMD ["sh", "-c", "exec uvicorn server:app --host 0.0.0.0 --port 5001 --workers \"${PII_WORKERS:-1}\""] - # Pinned spaCy models (en + es/it/pl/fi, ~2.2GB total). Downloaded with # retries/resume — the large wheels truncate on flaky networks if pip fetches -# the URLs directly. Shared by every terminal image so PII_ENGINE=spacy works -# everywhere (the gliner opt-in stays reversible without an image swap). -FROM base AS spacy-models +# the URLs directly. ARG SPACY_MODELS="en_core_web_lg-3.8.0 es_core_news_lg-3.8.0 it_core_news_lg-3.8.0 pl_core_news_lg-3.8.0 fi_core_news_lg-3.8.0" RUN --mount=type=cache,target=/root/.cache/pip \ for model in ${SPACY_MODELS}; do \ @@ -69,15 +43,14 @@ RUN --mount=type=cache,target=/root/.cache/pip \ pip install /tmp/*.whl && \ rm /tmp/*.whl -# --- GLiNER (CPU) ------------------------------------------------------------ -FROM spacy-models AS gliner - -# torch pinned here (not requirements-gliner.txt) because the CPU and CUDA -# targets install the same version from different wheel indexes. 2.11.0 is the +# --- GLiNER engine deps ------------------------------------------------------- +# torch is pinned here (not requirements-gliner.txt) because the CPU and CUDA +# builds install the same version from different wheel indexes. 2.11.0 is the # newest release published on both the cpu and cu128 indexes for py312. ARG TORCH_VERSION=2.11.0 +ARG TORCH_INDEX_URL=https://download.pytorch.org/whl/cpu RUN --mount=type=cache,target=/root/.cache/pip \ - pip install torch==${TORCH_VERSION} --index-url https://download.pytorch.org/whl/cpu + pip install torch==${TORCH_VERSION} --index-url ${TORCH_INDEX_URL} COPY apps/pii/requirements-gliner.txt ./requirements-gliner.txt RUN --mount=type=cache,target=/root/.cache/pip \ @@ -105,54 +78,38 @@ RUN python -c "from gliner import GLiNER; GLiNER.from_pretrained('${GLINER_MODEL chmod -R a+rX /opt/hf-cache ENV HF_HUB_OFFLINE=1 -# Bench + tests ride along only in this image — it's the only one with both -# engines installed (see apps/pii/scripts/bench_engines.py). pytest/httpx are -# baked in so the documented `docker run ... python -m pytest tests` works -# as-is (the runtime user has no writable HOME for pip install --user). +# pytest/httpx for the in-image test suites (tests/) — baked in because the +# runtime user has no writable HOME for pip install --user. COPY apps/pii/requirements-dev.txt ./requirements-dev.txt RUN --mount=type=cache,target=/root/.cache/pip \ pip install -r requirements-dev.txt +RUN groupadd -g 1001 pii && \ + useradd -u 1001 -g pii pii && \ + chown -R pii:pii /app + COPY --chown=pii:pii apps/pii/server.py apps/pii/engines.py ./ COPY --chown=pii:pii apps/pii/scripts ./scripts COPY --chown=pii:pii apps/pii/tests ./tests USER pii -# --- GLiNER (CUDA) scaffold — wired for the GPU fleet follow-up -------------- -FROM spacy-models AS gliner-gpu - -ARG TORCH_VERSION=2.11.0 -RUN --mount=type=cache,target=/root/.cache/pip \ - pip install torch==${TORCH_VERSION} --index-url https://download.pytorch.org/whl/cu128 - -COPY apps/pii/requirements-gliner.txt ./requirements-gliner.txt -RUN --mount=type=cache,target=/root/.cache/pip \ - pip install -r requirements-gliner.txt - -ARG SPACY_SM_MODELS="en_core_web_sm-3.8.0 es_core_news_sm-3.8.0 it_core_news_sm-3.8.0 pl_core_news_sm-3.8.0 fi_core_news_sm-3.8.0" -RUN --mount=type=cache,target=/root/.cache/pip \ - for model in ${SPACY_SM_MODELS}; do \ - whl="${model}-py3-none-any.whl"; \ - curl -fL --retry 5 --retry-delay 5 --retry-all-errors -C - \ - -o "/tmp/${whl}" \ - "https://github.com/explosion/spacy-models/releases/download/${model}/${whl}" || exit 1; \ - done && \ - pip install /tmp/*.whl && \ - rm /tmp/*.whl - -ENV HF_HOME=/opt/hf-cache -ARG GLINER_MODEL=urchade/gliner_multi_pii-v1 -RUN python -c "from gliner import GLiNER; GLiNER.from_pretrained('${GLINER_MODEL}')" && \ - chmod -R a+rX /opt/hf-cache -ENV HF_HUB_OFFLINE=1 - -COPY --chown=pii:pii apps/pii/server.py apps/pii/engines.py ./ -COPY --chown=pii:pii apps/pii/scripts ./scripts +# Listen on 5001. Runs as its own ECS service (separate task), reached via PII_URL; +# 5001 avoids colliding with the app's 3000 in local/compose runs on one host. +EXPOSE 5001 -USER pii +# start-period covers the model cold start. With PII_WORKERS>1 each worker loads +# the five spaCy models independently and in parallel, so allow generous headroom +# (memory-bandwidth contention stretches the wall-time beyond the single-worker case). +HEALTHCHECK --interval=30s --timeout=5s --start-period=300s --retries=3 \ + CMD curl -fsS http://localhost:5001/health || exit 1 -# --- DEFAULT (last stage): the lean spaCy image, content-equivalent to today - -FROM spacy-models AS spacy -COPY --chown=pii:pii apps/pii/server.py apps/pii/engines.py ./ -USER pii +# 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. +# NB for the gliner engine: EACH worker loads its own GLiNER model copy (into GPU +# memory when on cuda), so GPU deployments generally want PII_WORKERS=1 per GPU. +CMD ["sh", "-c", "exec uvicorn server:app --host 0.0.0.0 --port 5001 --workers \"${PII_WORKERS:-1}\""] diff --git a/helm/sim/values.schema.json b/helm/sim/values.schema.json index 71ec431f6d5..37f518542b2 100644 --- a/helm/sim/values.schema.json +++ b/helm/sim/values.schema.json @@ -714,7 +714,7 @@ "engine": { "type": "string", "enum": ["spacy", "gliner"], - "description": "NER engine; gliner requires the -gliner image variant" + "description": "NER engine (spacy default, gliner opt-in; both ship in the image)" }, "device": { "type": "string", diff --git a/helm/sim/values.yaml b/helm/sim/values.yaml index ee62dee233a..a12713a5bec 100644 --- a/helm/sim/values.yaml +++ b/helm/sim/values.yaml @@ -824,10 +824,9 @@ pii: # NER engine: "spacy" (default) or "gliner" (opt-in zero-shot transformer NER # for PERSON/LOCATION/NRP/DATE_TIME; regex/checksum recognizers are identical - # on both engines). gliner requires the gliner image variant — the `-gliner` - # published tag (e.g. image.tag: latest-gliner) or a local - # `docker build --target gliner -f docker/pii.Dockerfile` — and fails fast at - # startup on the default image. + # on both engines). The image ships both engines, so this is a pure env flip + # — no image change needed. Note: gliner on CPU is orders of magnitude slower + # than spacy; it is intended for GPU nodes. engine: "spacy" # Torch device for the gliner engine: "cpu", "cuda", or "cuda:N". Empty