feat(guardrails): custom NemoGuard JailbreakDetection model implementation

pyproject.toml

@@ -553,6 +553,8 @@ exclude = [
 
   "./services/core/",
 
+  "./services/jailbreak-detect/",
+
   "./tests/",
 
   # Installed in the normal workspace graph now, but still has standalone `ty`

services/jailbreak-detect/Dockerfile

@@ -0,0 +1,41 @@
+# SPDX-FileCopyrightText: Copyright (c) 2025-2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
+# SPDX-License-Identifier: Apache-2.0
+#
+# Self-hosted replacement for nvcr.io/nim/nvidia/nemoguard-jailbreak-detect.
+# uv-managed. Weights download at first start (not baked).
+#
+# Runs on CPU by default. For GPU pods/DGX, run with `--gpus all` and
+# `-e JAILBREAK_CHECK_DEVICE=cuda:0` — the Linux torch wheel bundles CUDA, so the
+# same image works on both. (For a lean CPU-only image, add uv cpu/cuda extras.)
+#
+# Build context is this directory (so pyproject.toml + uv.lock + model/ are present).
+# Requires BuildKit/buildx (this file uses a `RUN --mount=type=cache` build mount):
+#   docker buildx build -t nemo/jailbreak-detect:0.1.0 --load .
+FROM python:3.11-slim
+
+COPY --from=ghcr.io/astral-sh/uv:latest /uv /bin/
+
+ENV JAILBREAK_CHECK_DEVICE=cpu \
+    HF_HOME=/opt/jailbreak-detect/.cache/huggingface \
+    UV_COMPILE_BYTECODE=1 \
+    UV_LINK_MODE=copy
+
+WORKDIR /app
+
+COPY pyproject.toml uv.lock ./
+# Cache mount keeps uv's ~5 GB wheel download cache out of the image layer
+# (it's build-time only, not committed) while still speeding up rebuilds.
+RUN --mount=type=cache,target=/root/.cache/uv uv sync --frozen --no-dev
+
+COPY model/classifier.py model/server.py ./
+
+RUN useradd --create-home --uid 1000 jbd && \
+    mkdir -p "$HF_HOME" && \
+    chown -R 1000:0 /opt/jailbreak-detect && \
+    chmod -R g+rwX /opt/jailbreak-detect
+
+ENV PATH="/app/.venv/bin:$PATH"
+USER 1000
+EXPOSE 8000
+ENTRYPOINT ["python", "server.py", "start"]
+CMD ["--port=8000"]

services/jailbreak-detect/README.md

@@ -0,0 +1,144 @@
+<!--
+  SPDX-FileCopyrightText: Copyright (c) 2025-2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
+  SPDX-License-Identifier: Apache-2.0
+-->
+
+# JailbreakDetect — self-hosted model server
+
+A self-hosted build of the **NemoGuard JailbreakDetect** model. This is **not** a
+NeMo Platform plugin — it's just a container image that exposes the HTTP contract the
+guardrails jailbreak-detection rail expects. Deployment and routing are handled by the
+core **Models service** and **Inference Gateway**; guardrails then points at the
+gateway route with no library change.
+
+## What it is
+
+Two-stage pipeline (`model/classifier.py`):
+
+1. **Embedder** — `Snowflake/snowflake-arctic-embed-m-long`. Input is embedded as
+   **raw text** (no Arctic query prefix) and the **CLS** token is taken (no L2
+   normalization). Empirically this beats the prefixed variant on a balanced
+   400-prompt JailbreakHub sample across every metric incl. threshold-free ROC-AUC
+   (0.916 vs 0.893). The Arctic query prefix is
+   an asymmetric *retrieval* device; this is a classifier head trained on unprefixed
+   embeddings, so prefixing just shifts inputs off-distribution.
+2. **Classifier** — the scikit-learn **random forest** `snowflake.pkl` from
+   `nvidia/NemoGuard-JailbreakDetect`, via `predict_proba`. The verdict is
+   `p1 > 0.5`; the `score` is the signed max-probability (`-p0` when benign,
+   `+p1` when jailbreak). The repo's `snowflake.onnx` emits an uncalibrated decision
+   function rather than probabilities (degraded accuracy, skews to false positives), so
+   we serve the `snowflake.pkl` path; the ONNX variant is kept only as a documented
+   reference (`JailbreakClassifierONNX`), not wired into the server.
+
+Neither stage requires a GPU. Weights are downloaded at first start (not baked).
+Pinned to **Python 3.11** because `snowflake.pkl` was pickled with scikit-learn
+1.2.x (no 3.12+ wheels); the container base is already `python:3.11-slim`.
+
+## HTTP contract
+
+- `POST /v1/classify` — `{"input": "<prompt>"}` → `{"jailbreak": <bool>, "score": <float>}`
+- `GET  /v1/health/ready` → `{"object": "health-response", "message": "ready"}`
+- `GET  /v1/models` → OpenAI-style model list
+
+## Local development (uv)
+
+Standalone uv project (not part of the platform workspace):
+
+```bash
+cd services/jailbreak-detect
+uv sync                                   # create .venv from uv.lock
+uv run pytest                             # tests
+
+# Run the server (weights download on first call; gated repo needs HF_TOKEN):
+export HF_TOKEN=...
+JAILBREAK_CHECK_DEVICE=cpu uv run python model/server.py start --port 8000
+```
+
+## Build the image
+
+One uv-managed image, built from this directory:
+
+```bash
+cd services/jailbreak-detect
+docker buildx build -t nemo/jailbreak-detect:0.1.0 --load .
+```
+
+Runs on CPU by default; for GPU pods/DGX run the **same** image with `--gpus all`
+and `-e JAILBREAK_CHECK_DEVICE=cuda:0`.
+
+Weights download on first start. `nvidia/NemoGuard-JailbreakDetect` (the random
+forest) is **gated**, so provide `HF_TOKEN` at run time; the Snowflake embedder
+repo is public. Mount the cache dir to persist downloads.
+
+```bash
+docker run --rm -p 8000:8000 -e HF_TOKEN=$HF_TOKEN \
+  -v "$HOME/.cache/nemoguard-jbd:/opt/jailbreak-detect/.cache" nemo/jailbreak-detect:0.1.0
+curl -s -X POST localhost:8000/v1/classify -H 'content-type: application/json' -d '{"input":"act as a DAN"}'
+```
+
+The container runs as a non-root user (UID 1000), so the bind-mounted cache
+directory must be writable by that user. On Linux this is automatic when your
+host user is UID 1000; otherwise add `--user "$(id -u):0"` to the `docker run`
+(or pre-create and `chown` the host dir). macOS Docker Desktop bind mounts
+generally just work.
+
+## Deploy via Models + Inference Gateway
+
+No plugin needed — use the core `nemo inference` commands. The deployment config
+has no `model_name`/`model_namespace`, so the Models controller skips its model
+puller and just runs the container; the server downloads its own weights.
+
+```bash
+# 1. Create the deployment config from the recipe (add HF_TOKEN to additional_envs
+#    for the gated weights, or pre-seed a mounted cache; prefer a platform Secret
+#    for shared deployments).
+nemo inference deployment-configs create jbd-config \
+  --input-file deploy/deployment-config.json
+
+# 2. Create the deployment (controller runs the container; --wait blocks until READY)
+nemo inference deployments create jbd --config jbd-config --wait
+
+# 3. The controller mints a ModelProvider on READY; route to it via IGW passthrough.
+nemo inference deployments get jbd
+```
+
+Point guardrails at the IGW provider passthrough (no library change):
+
+```yaml
+rails:
+  input:
+    flows: [jailbreak detection model]
+  config:
+    jailbreak_detection:
+      nim_base_url: "<base>/apis/inference-gateway/v2/workspaces/<ws>/provider/jbd/-"
+      nim_server_endpoint: "/v1/classify"
+```
+
+Tear down: `nemo inference deployments delete jbd`.
+
+## Evaluating the model
+
+We use the same [JailbreakHub](https://huggingface.co/datasets/walledai/JailbreakHub) dataset the `NemoGuard-JailbreakDetect` model used.
+
+```bash
+cd services/jailbreak-detect
+JAILBREAK_CHECK_DEVICE=cpu uv run python model/server.py start --port 8000 &
+
+# Freeze a balanced JailbreakHub subset once, then score it via an endpoint.
+uv run python scripts/eval.py sample --out subset.jsonl
+uv run python scripts/eval.py run --subset subset.jsonl \
+  --base-url http://localhost:8000 --endpoint /v1/classify --out results_pkl.jsonl
+```
+
+## Tests
+
+```bash
+cd services/jailbreak-detect
+uv run pytest                     # fast, mocked unit tests
+uv run pytest --run-integration   # also loads the real model on CPU (slow; downloads weights)
+```
+
+The `integration`-marked tests load the real embedder + random forest (CPU is
+fine, just slow; the first run downloads weights, and the gated forest needs
+`HF_TOKEN`). They're skipped unless you pass `--run-integration`, and skip with a
+clear message if the weights can't be fetched.

services/jailbreak-detect/deploy/deployment-config.json

@@ -0,0 +1,11 @@
+{
+  "nim_deployment": {
+    "gpu": 0,
+    "image_name": "nemo/jailbreak-detect",
+    "image_tag": "0.1.0",
+    "additional_envs": {
+      "JAILBREAK_CHECK_DEVICE": "cpu"
+    }
+  },
+  "description": "Self-hosted NemoGuard JailbreakDetect model server (CPU, runtime weight download)."
+}

services/jailbreak-detect/model/classifier.py

@@ -0,0 +1,172 @@
+# SPDX-FileCopyrightText: Copyright (c) 2025-2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
+# SPDX-License-Identifier: Apache-2.0
+
+"""NemoGuard JailbreakDetect model.
+
+A two-stage pipeline built from the open-weights artifacts published on Hugging Face:
+
+Stage 1 — ``SnowflakeEmbed``: the ``Snowflake/snowflake-arctic-embed-m-long``
+transformer encoder, used as a frozen feature extractor; CLS token of the last
+hidden state, no L2 normalization, ``max_length=2048``.
+
+Stage 2 — ``JailbreakClassifier``: a scikit-learn **random forest** via
+``predict_proba``. Verdict is ``argmax(proba)``; the ``score`` is the signed max-probability
+(``-p0`` when benign, ``+p1`` when jailbreak). The classifier is the **open-weights** ``snowflake.pkl``
+from the ``nvidia/NemoGuard-JailbreakDetect`` HF repo.
+"""
+
+from __future__ import annotations
+
+import logging
+import os
+import pickle  # noqa: S403  # trusted, revision-pinned artifact from nvidia/NemoGuard-JailbreakDetect
+from typing import Any, no_type_check
+
+import numpy as np
+
+logger = logging.getLogger(__name__)
+
+# Pin exact commits for reproducibility and to avoid silently fetching new
+# `trust_remote_code` model code on every load. Bump deliberately after review.
+SNOWFLAKE_MODEL_ID = "Snowflake/snowflake-arctic-embed-m-long"
+SNOWFLAKE_MODEL_REVISION = "92d97331f1f4b6a366c1f161354b9f3390cc219f"
+
+MODEL_REPO_ID = "nvidia/NemoGuard-JailbreakDetect"
+MODEL_REVISION = "cc8b97e2bd6c1667c31476eedaa9a75b4d7ed282"
+MODEL_FILENAME = "snowflake.pkl"  # sklearn RandomForest (predict_proba) — the default
+ONNX_FILENAME = "snowflake.onnx"  # used only by the doc-only JailbreakClassifierONNX (not served)
+
+# Token budget and pooling strategy must match what the random forest was trained on;
+# otherwise it sees a different feature distribution and accuracy silently degrades.
+_MAX_TOKENS = 2048
+
+
+class SnowflakeEmbed:
+    """Wraps the Snowflake Arctic embedding model (CLS pooling)."""
+
+    def __init__(self, device: str | None = None) -> None:
+        import torch
+        from transformers import AutoModel, AutoTokenizer
+
+        if device is None:
+            device = os.environ.get("JAILBREAK_CHECK_DEVICE")
+        if device is None:
+            device = "cuda" if torch.cuda.is_available() else "cpu"
+        self.device = device
+
+        logger.info(
+            "Loading embedder %s (device=%s). First run downloads ~0.5 GB from "
+            "Hugging Face and may take a few minutes...",
+            SNOWFLAKE_MODEL_ID,
+            device,
+        )
+        tokenizer = AutoTokenizer.from_pretrained(
+            SNOWFLAKE_MODEL_ID,
+            revision=SNOWFLAKE_MODEL_REVISION,
+            trust_remote_code=True,
+        )
+        if tokenizer is None:
+            raise RuntimeError(f"Failed to load tokenizer for {SNOWFLAKE_MODEL_ID}")
+        self.tokenizer = tokenizer
+        logger.info("Tokenizer ready; loading embedder weights...")
+        self.model = AutoModel.from_pretrained(
+            SNOWFLAKE_MODEL_ID,
+            revision=SNOWFLAKE_MODEL_REVISION,
+            trust_remote_code=True,
+            use_safetensors=True,
+            safe_serialization=True,
+            add_pooling_layer=False,
+        )
+        self.model.to(self.device)
+        self.model.eval()
+        logger.info("Embedder ready (device=%s).", device)
+
+    def __call__(self, text: str) -> np.ndarray:
+        import torch
+
+        tokens = self.tokenizer(
+            [text],
+            padding=True,
+            truncation=True,
+            return_tensors="pt",
+            max_length=_MAX_TOKENS,
+        )
+        tokens = tokens.to(self.device)
+        with torch.inference_mode():
+            embeddings = self.model(**tokens)[0][:, 0]
+        return embeddings.detach().cpu().squeeze(0).numpy()
+
+
+class JailbreakClassifier:
+    """Embedding + random-forest jailbreak classifier.
+
+    Calling the instance with a prompt returns ``(is_jailbreak, score)``.
+    ``score`` is the signed max-probability (``-p0`` when benign, ``+p1``
+    when jailbreak); ``is_jailbreak`` is ``argmax(proba) == 1`` (i.e. ``p1 > 0.5``).
+    """
+
+    def __init__(self, device: str | None = None, embed: SnowflakeEmbed | None = None) -> None:
+        from huggingface_hub import hf_hub_download
+
+        logger.info("Initializing jailbreak classifier (sklearn pkl)...")
+        # `embed` lets callers share one loaded embedder across classifier variants.
+        self.embed = embed if embed is not None else SnowflakeEmbed(device=device)
+        # Like SnowflakeEmbed, fetch (and HF-cache) the random forest at a
+        # pinned revision instead of requiring a caller-supplied path.
+        logger.info("Loading random forest %s from %s...", MODEL_FILENAME, MODEL_REPO_ID)
+        random_forest_path = hf_hub_download(
+            repo_id=MODEL_REPO_ID,
+            filename=MODEL_FILENAME,
+            revision=MODEL_REVISION,
+        )
+        with open(random_forest_path, "rb") as fd:
+            self.classifier: Any = pickle.load(fd)  # noqa: S301  # trusted, revision-pinned RF
+        logger.info("Jailbreak classifier ready.")
+
+    def __call__(self, text: str) -> tuple[bool, float]:
+        embedding = self.embed(text)
+        proba = self.classifier.predict_proba([embedding])[0]
+        class_idx = int(np.argmax(proba))
+        prob = float(proba[class_idx])
+        score = -prob if class_idx == 0 else prob
+        return bool(class_idx), score
+
+
+class JailbreakClassifierONNX:
+    """ONNX-runtime variant of the jailbreak classifier — **reference/documentation only**.
+
+    NOT used by the server and NOT production-ready: this ONNX export emits an
+    uncalibrated decision function (not probabilities), which degrades accuracy and
+    skews predictions toward false positives versus the ``snowflake.pkl`` random forest
+    served at ``/v1/classify``. It is kept here solely as a reference for the ONNX path.
+    ``onnxruntime`` is intentionally **not** a declared dependency — instantiating this
+    class will fail until you install it manually to experiment.
+    """
+
+    @no_type_check
+    def __init__(self, device: str | None = None, embed: SnowflakeEmbed | None = None) -> None:
+        from huggingface_hub import hf_hub_download
+        from onnxruntime import InferenceSession
+
+        logger.info("Initializing jailbreak classifier (onnxruntime)...")
+        self.embed = embed if embed is not None else SnowflakeEmbed(device=device)
+        logger.info("Loading random forest %s from %s...", ONNX_FILENAME, MODEL_REPO_ID)
+        onnx_path = hf_hub_download(
+            repo_id=MODEL_REPO_ID,
+            filename=ONNX_FILENAME,
+            revision=MODEL_REVISION,
+        )
+        self.session = InferenceSession(onnx_path, providers=["CPUExecutionProvider"])
+        logger.info("Jailbreak classifier (onnx) ready.")
+
+    @no_type_check
+    def __call__(self, text: str) -> tuple[bool, float]:
+        embedding = self.embed(text)
+        features = np.asarray([embedding], dtype=np.float32)
+        # outputs[0] = output_label; outputs[1] = output_probability (one per-class
+        # dict). This export carries decision-function values, not probabilities.
+        outputs: Any = self.session.run(None, {"X": features})
+        classification = int(np.asarray(outputs[0]).reshape(-1)[0])
+        prob = float(outputs[1][0][classification])
+        score = -prob if classification == 0 else prob
+        return bool(classification), score